def test_set_boost_factor_with_defaults():
fake_region = 'westeros-1'
fake_pool = 'default'
base_path = cluster_boost.get_zk_boost_path(fake_region, fake_pool)
with patch_zk_client() as mock_zk_client:
cluster_boost.set_boost_factor(fake_region, fake_pool)
expected_end_time = float(TEST_CURRENT_TIME.timestamp()
) + 60 * cluster_boost.DEFAULT_BOOST_DURATION
assert mock_zk_client.set.call_args_list == [
mock.call(
base_path + '/end_time',
str(expected_end_time).encode('utf-8'),
),
mock.call(
base_path + '/factor',
str(cluster_boost.DEFAULT_BOOST_FACTOR).encode('utf-8'),
),
mock.call(
base_path + '/expected_load',
'0'.encode('utf-8'),
),
]
def test_set_boost_factor_with_active_boost_override():
fake_region = 'westeros-1'
fake_pool = 'default'
base_path = cluster_boost.get_zk_boost_path(fake_region, fake_pool)
fake_end_time = float(TEST_CURRENT_TIME.timestamp()) + 10
fake_boost_factor = 1.5
fake_expected_load = 80
mock_boost_values = {
base_path + '/end_time': str(fake_end_time).encode('utf-8'),
base_path + '/factor': str(fake_boost_factor).encode('utf-8'),
base_path + '/expected_load': str(fake_expected_load).encode('utf-8'),
}
# patch zk client so that it returns an end time that
# indicates an active boost
with patch_zk_client(mock_boost_values) as mock_zk_client:
# we need the zk.set to actually override the initial mocked values
def mock_set(key, value):
mock_boost_values[key] = value
mock_zk_client.set = mock_set
# set boost will go through with an active boost if override is toggled on
assert cluster_boost.set_boost_factor(
region=fake_region,
pool=fake_pool,
override=True,
)
def paasta_cluster_boost():
""" Set, Get or clear a boost on a paasta cluster for a given pool in a given region
:returns: None
"""
system_config = load_system_paasta_config()
args = parse_args()
if args.verbose >= 2:
logging.basicConfig(level=logging.DEBUG)
elif args.verbose == 1:
logging.basicConfig(level=logging.INFO)
else:
logging.basicConfig(level=logging.WARNING)
action = args.action
pool = args.pool
if not system_config.get_cluster_boost_enabled():
paasta_print('ERROR: cluster_boost feature is not enabled.')
return False
regions = get_regions(pool)
if len(regions) == 0:
paasta_print('ERROR: no slaves found in pool {}'.format(pool))
return False
for region in regions:
if action == 'set':
if not cluster_boost.set_boost_factor(
region=region,
pool=pool,
factor=args.boost,
duration_minutes=args.duration,
override=args.override,
):
paasta_print('ERROR: Failed to set the boost for pool {}, region {}.'.format(pool, region))
return False
elif action == 'status':
pass
elif action == 'clear':
if not cluster_boost.clear_boost(pool=pool, region=region):
paasta_print('ERROR: Failed to clear the boost for pool {}, region {}.')
return False
else:
raise NotImplementedError("Action: '%s' is not implemented." % action)
return False
paasta_print('Current boost value for pool: {}, region: {}: {}'.format(
pool, region, cluster_boost.get_boost_factor(
region=region, pool=pool,
),
))
def paasta_cluster_boost():
""" Set, Get or clear a boost on a paasta cluster for a given pool in a given region
:returns: None
"""
args = parse_args()
if args.verbose >= 2:
logging.basicConfig(level=logging.DEBUG)
elif args.verbose == 1:
logging.basicConfig(level=logging.INFO)
else:
logging.basicConfig(level=logging.WARNING)
action = args.action
region = args.region
pool = args.pool
if action == 'set':
# Let's disallow people to set boost on a non existing pool
if not check_pool_exist(pool=pool, region=region):
paasta_print('Could not find the pool {} in the region {}'.format(
pool, region))
return False
if not cluster_boost.set_boost_factor(
region=region,
pool=pool,
factor=args.boost,
duration_minutes=args.duration,
override=args.override,
):
paasta_print('Failed! Check the logs.')
return False
elif action == 'get':
paasta_print('Current boost value: {}'.format(
cluster_boost.get_boosted_load(
region=region,
pool=pool,
current_load=1.0,
)))
elif action == 'clear':
if not cluster_boost.clear_boost(pool=pool, region=region):
paasta_print('Failed! Check the logs.')
return False
else:
raise NotImplementedError("Action: '%s' is not implemented." % action)
return False
def test_set_boost_factor_with_active_boost():
fake_region = 'westeros-1'
fake_pool = 'default'
base_path = cluster_boost.get_zk_boost_path(fake_region, fake_pool)
fake_end_time = float(TEST_CURRENT_TIME.timestamp()) + 10
fake_boost_factor = 1.5
fake_expected_load = 80
# patch zk client so that it returns an end time that
# indicates an active boost
with patch_zk_client({
base_path + '/end_time':
str(fake_end_time).encode('utf-8'),
base_path + '/factor':
str(fake_boost_factor).encode('utf-8'),
base_path + '/expected_load':
str(fake_expected_load).encode('utf-8'),
}):
# by default, set boost should not go through if there's an active boost
assert not cluster_boost.set_boost_factor(region=fake_region,
pool=fake_pool)