import datetime from core.throttling_utils import ( ThrottleManager, CooldownThrottlingRule, WorkingHoursThrottlingRule, DynamicThrottlingRule, ) def test_working_hours_throttling_rule_checks(): working_hours_rule = WorkingHoursThrottlingRule( working_hours={ "start": datetime.datetime.now().time(), "end": (datetime.datetime.now() + datetime.timedelta(hours=1)).time(), } ) assert working_hours_rule() == True def test_working_hours_throttling_rule_does_not_check(): working_hours_rule = WorkingHoursThrottlingRule( working_hours={ "start": (datetime.datetime.now() + datetime.timedelta(hours=1)).time(), "end": (datetime.datetime.now() + datetime.timedelta(hours=2)).time(), } ) assert working_hours_rule() == False def test_cooldown_throttling_rule_checks(): time_generator = lambda: 60 cooldown_rule = CooldownThrottlingRule(cooldown_time_generator=time_generator) assert ( cooldown_rule( last_attempt_timestamp=datetime.datetime.now() + datetime.timedelta(seconds=-120) ) == True ) def test_cooldown_throttling_rule_does_not_check(): time_generator = lambda: 60 cooldown_rule = CooldownThrottlingRule(cooldown_time_generator=time_generator) assert cooldown_rule(last_attempt_timestamp=datetime.datetime.now()) == False def test_dynamic_rule_checks(): mock_check = lambda: True rule = DynamicThrottlingRule(any_callable=mock_check) assert rule() == True def test_dynamic_rule_does_not_check(): mock_check = lambda: False rule = DynamicThrottlingRule(any_callable=mock_check) assert rule() == False def test_dynamic_rule_arguments_pass_properly(): def pass_a_bool(some_bool): return some_bool rule = DynamicThrottlingRule(pass_a_bool) assert (rule(some_bool=True) == True) and (rule(some_bool=False) == False) def test_throttle_manager_checks_rules(): throttle_manager = ThrottleManager() def pass_a_bool(some_bool): return some_bool some_rules = [ WorkingHoursThrottlingRule( working_hours={ "start": datetime.datetime.now().time(), "end": (datetime.datetime.now() + datetime.timedelta(hours=1)).time(), } ), CooldownThrottlingRule(cooldown_time_generator=lambda: 0), DynamicThrottlingRule(any_callable=pass_a_bool), ] some_arguments = [[], ["last_attempt_timestamp"], ["some_bool"]] some_rules_and_arguments = zip(some_rules, some_arguments) for rule, arguments in some_rules_and_arguments: throttle_manager.add_rule(rule, required_argument_names=arguments) assert throttle_manager.allow_next_task( last_attempt_timestamp=datetime.datetime.now(), some_bool=True )