def test_send_clusterman_metrics(mock_load_system_paasta_config,
mock_clusterman_metrics):
fake_region = "westeros-1"
fake_pool = "default"
base_path = load_boost.get_zk_cluster_boost_path(fake_region, fake_pool)
mock_load_system_paasta_config.return_value.get_cluster.return_value = (
"westeros-prod")
mock_clusterman_metrics.generate_key_with_dimensions = (
lambda key, dims: f"{key}|{dims}")
mock_writer = (mock_clusterman_metrics.ClustermanMetricsBotoClient().
get_writer().__enter__())
with patch_zk_client():
load_boost.set_boost_factor(
zk_boost_path=base_path,
region=fake_region,
pool=fake_pool,
factor=1.3,
duration_minutes=10,
)
expected_metrics_dimensions = {
"cluster": "westeros-prod",
"pool": "default"
}
expected_metrics_key = f"boost_factor|{expected_metrics_dimensions}"
assert mock_writer.send.call_args_list == [
mock.call((expected_metrics_key, TEST_CURRENT_TIME.timestamp(), 1.3)),
mock.call((expected_metrics_key,
TEST_CURRENT_TIME.timestamp() + 10 * 60, 1.0)),
]
def test_send_clusterman_metrics(mock_load_system_paasta_config, mock_clusterman_metrics):
fake_region = 'westeros-1'
fake_pool = 'default'
base_path = load_boost.get_zk_cluster_boost_path(fake_region, fake_pool)
mock_load_system_paasta_config.return_value.get_cluster.return_value = 'westeros-prod'
mock_clusterman_metrics.generate_key_with_dimensions = lambda key, dims: f'{key}|{dims}'
mock_writer = mock_clusterman_metrics.ClustermanMetricsBotoClient().get_writer().__enter__()
with patch_zk_client():
load_boost.set_boost_factor(
zk_boost_path=base_path,
region=fake_region,
pool=fake_pool,
factor=1.3,
duration_minutes=10,
)
expected_metrics_dimensions = {'cluster': 'westeros-prod', 'pool': 'default'}
expected_metrics_key = f'boost_factor|{expected_metrics_dimensions}'
assert mock_writer.send.call_args_list == [
mock.call((expected_metrics_key, TEST_CURRENT_TIME.timestamp(), 1.3)),
mock.call((expected_metrics_key, TEST_CURRENT_TIME.timestamp() + 10 * 60, 1.0)),
]
def test_set_boost_factor_with_active_boost_override():
fake_region = "westeros-1"
fake_pool = "default"
base_path = load_boost.get_zk_cluster_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 load_boost.set_boost_factor(
zk_boost_path=
f"/paasta_cluster_autoscaler/{fake_region}/{fake_pool}/boost",
override=True,
)
def test_set_boost_factor_with_defaults():
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.set_boost_factor(base_path)
expected_end_time = (float(TEST_CURRENT_TIME.timestamp()) +
60 * load_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(load_boost.DEFAULT_BOOST_FACTOR).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_set_boost_factor_with_defaults():
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.set_boost_factor(base_path)
expected_end_time = float(TEST_CURRENT_TIME.timestamp()) + 60 * load_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(load_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():
fake_region = 'westeros-1'
fake_pool = 'default'
base_path = load_boost.get_zk_cluster_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 load_boost.set_boost_factor(zk_boost_path=base_path)
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
