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)