Repository: Chia-Network/chia-blockchain
File: chia/_tests/core/server/test_rate_limits.py (Python)

Code before:
# * number of bytes / 60 seconds for transaction messages


@pytest.mark.anyio
@boolean_datacases(name="incoming", true="incoming", false="outgoing")
@boolean_datacases(name="tx_msg", true="tx", false="non-tx")
@boolean_datacases(name="limit_size", true="size-limit", false="count-limit")
async def test_limits_v2(incoming: bool, tx_msg: bool, limit_size: bool, monkeypatch: pytest.MonkeyPatch) -> None:
    # this test uses a single message type, and alters the rate limit settings
    # for it to hit the different cases

    count = 1000
    message_data = b"\0" * 1024
    msg_type = ProtocolMessageTypes.new_transaction

    limits: dict[ProtocolMessageTypes, Union[RLSettings, Unlimited]]

    if limit_size:
        agg_limit = RLSettings(False, count * 2, 1024, count * len(message_data))
    else:
        agg_limit = RLSettings(False, count, 1024, count * 2 * len(message_data))

    if limit_size:
        limits = {msg_type: RLSettings(not tx_msg, count * 2, 1024, count * len(message_data))}
    else:
        limits = {msg_type: RLSettings(not tx_msg, count, 1024, count * 2 * len(message_data))}

    def mock_get_limits(
        our_capabilities: list[Capability], peer_capabilities: list[Capability]
    ) -> tuple[dict[ProtocolMessageTypes, Union[RLSettings, Unlimited]], RLSettings]:
        return limits, agg_limit

    import chia.server.rate_limits

    monkeypatch.setattr(chia.server.rate_limits, "get_rate_limits_to_use", mock_get_limits)

    r = RateLimiter(incoming=incoming, get_time=lambda: 0)
    msg = make_msg(msg_type, message_data)

    for i in range(count):
        assert r.process_msg_and_check(msg, rl_v2, rl_v2) is None

    expected_msg = ""

    if limit_size:
        if not tx_msg:
            expected_msg += "non-tx size:"
        else:
            expected_msg += "cumulative size:"
        expected_msg += f" {(count + 1) * len(message_data)} > {count * len(message_data) * 1.0}"
    else:

Review comment (suggestion):
is the intent of these 1024s to be `len(message_data)`?  if so, might be worth changing the four cases here.

Code after:
# * number of bytes / 60 seconds for transaction messages


@pytest.mark.anyio
@boolean_datacases(name="incoming", true="incoming", false="outgoing")
@boolean_datacases(name="tx_msg", true="tx", false="non-tx")
@boolean_datacases(name="limit_size", true="size-limit", false="count-limit")
async def test_limits_v2(incoming: bool, tx_msg: bool, limit_size: bool, monkeypatch: pytest.MonkeyPatch) -> None:
    # this test uses a single message type, and alters the rate limit settings
    # for it to hit the different cases

    count = 1000
    message_data = b"\0" * 1024
    msg_type = ProtocolMessageTypes.new_transaction

    limits: dict[ProtocolMessageTypes, Union[RLSettings, Unlimited]]

    if limit_size:
        agg_limit = RLSettings(False, count * 2, len(message_data), count * len(message_data))
    else:
        agg_limit = RLSettings(False, count, len(message_data), count * 2 * len(message_data))

    if limit_size:
        limits = {msg_type: RLSettings(not tx_msg, count * 2, len(message_data), count * len(message_data))}
    else:
        limits = {msg_type: RLSettings(not tx_msg, count, len(message_data), count * 2 * len(message_data))}

    def mock_get_limits(
        our_capabilities: list[Capability], peer_capabilities: list[Capability]
    ) -> tuple[dict[ProtocolMessageTypes, Union[RLSettings, Unlimited]], RLSettings]:
        return limits, agg_limit

    import chia.server.rate_limits

    monkeypatch.setattr(chia.server.rate_limits, "get_rate_limits_to_use", mock_get_limits)

    r = RateLimiter(incoming=incoming, get_time=lambda: 0)
    msg = make_msg(msg_type, message_data)

    for i in range(count):
        assert r.process_msg_and_check(msg, rl_v2, rl_v2) is None

    expected_msg = ""

    if limit_size:
        if not tx_msg:
            expected_msg += "non-tx size:"
        else:
            expected_msg += "cumulative size:"
        expected_msg += f" {(count + 1) * len(message_data)} > {count * len(message_data) * 1.0}"
    else: