def test_clear_boost():
fake_region = "westeros-1"
fake_pool = "default"
base_path = load_boost.get_zk_cluster_boost_path(fake_region, fake_pool)
with patch_zk_client() as mock_zk_client:
load_boost.clear_boost(base_path, region=fake_region, pool=fake_pool)
expected_end_time = float(TEST_CURRENT_TIME.timestamp())
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", "1".encode("utf-8")),
mock.call(base_path + "/expected_load", "0".encode("utf-8")),
]
def paasta_cluster_boost(action: str, pool: str, boost: float, duration: int,
override: bool) -> bool:
""" 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()
if not system_config.get_cluster_boost_enabled():
print("ERROR: cluster_boost feature is not enabled.")
return False
regions = system_config.get_boost_regions()
if len(regions) == 0:
print(
f"ERROR: no boost_regions configured in {system_config.directory}")
return False
for region in regions:
zk_boost_path = load_boost.get_zk_cluster_boost_path(region=region,
pool=pool)
if action == "set":
if not load_boost.set_boost_factor(
zk_boost_path=zk_boost_path,
region=region,
pool=pool,
factor=boost,
duration_minutes=duration,
override=override,
):
print(
f"ERROR: Failed to set the boost for pool {pool}, region {region}."
)
return False
elif action == "status":
pass
elif action == "clear":
if not load_boost.clear_boost(
zk_boost_path, region=region, pool=pool):
print(
"ERROR: Failed to clear the boost for pool {}, region {}.")
return False
else:
raise NotImplementedError("Action: '%s' is not implemented." %
action)
return False
print("Current boost value for path: {}: {}".format(
zk_boost_path,
load_boost.get_boost_factor(zk_boost_path=zk_boost_path)))
return True
def test_clear_boost():
fake_region = 'westeros-1'
fake_pool = 'default'
base_path = load_boost.get_zk_cluster_boost_path(fake_region, fake_pool)
with patch_zk_client() as mock_zk_client:
load_boost.clear_boost(base_path, region=fake_region, pool=fake_pool)
expected_end_time = float(TEST_CURRENT_TIME.timestamp())
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',
'1'.encode('utf-8'),
),
mock.call(
base_path + '/expected_load',
'0'.encode('utf-8'),
),
]
def paasta_cluster_boost(
action: str,
pool: str,
boost: float,
duration: int,
override: bool,
) -> bool:
""" 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()
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(f'ERROR: no slaves found in pool {pool}')
return False
for region in regions:
zk_boost_path = load_boost.get_zk_cluster_boost_path(
region=region,
pool=pool,
)
if action == 'set':
if not load_boost.set_boost_factor(
zk_boost_path=zk_boost_path,
region=region,
pool=pool,
factor=boost,
duration_minutes=duration,
override=override,
):
paasta_print(
f'ERROR: Failed to set the boost for pool {pool}, region {region}.'
)
return False
elif action == 'status':
pass
elif action == 'clear':
if not load_boost.clear_boost(
zk_boost_path,
region=region,
pool=pool,
):
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 path: {}: {}'.format(
zk_boost_path,
load_boost.get_boost_factor(zk_boost_path=zk_boost_path, ),
))
return True