drogon/tests/throttling_test.py

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() + datetime.timedelta(seconds=-5)).time(),
            "end": (datetime.datetime.now() + datetime.timedelta(seconds=5)).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() + datetime.timedelta(seconds=-5)
                ).time(),
                "end": (datetime.datetime.now() + datetime.timedelta(seconds=5)).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
    )