Repository: pydantic/pydantic-ai
File: tests/test_usage_limits.py (Python)

Code before:
        return 'result c'

    # Run with tool call limit of 6; expecting an error when trying to execute 3 more tools
    with pytest.raises(
        UsageLimitExceeded,
        match=re.escape('The next tool call(s) would exceed the tool_calls_limit of 6 (tool_calls=8).'),
    ):
        await agent.run('Use tools', usage_limits=UsageLimits(tool_calls_limit=6))

    # Only the first batch of 5 tools should have executed
    assert len(executed_tools) == 5


def test_usage_unknown_provider():
    assert RequestUsage.extract({}, provider='unknown', provider_url='', provider_fallback='') == RequestUsage()


def test_usage_limits_preserves_explicit_zero():
    """Test that explicit 0 token limits are preserved and not replaced by deprecated fallbacks."""
    # When new arg is 0 and deprecated arg is non-zero, should use 0 (not fallback)
    # Using type: ignore since we're intentionally testing mixed deprecated/new args
    limits = UsageLimits(input_tokens_limit=0, request_tokens_limit=123)  # type: ignore[call-overload]
    assert limits.input_tokens_limit == 0

    limits = UsageLimits(output_tokens_limit=0, response_tokens_limit=456)  # type: ignore[call-overload]
    assert limits.output_tokens_limit == 0

    # When new arg is not passed, should use deprecated fallback
    limits = UsageLimits(request_tokens_limit=123)
    assert limits.input_tokens_limit == 123

    limits = UsageLimits(response_tokens_limit=456)
    assert limits.output_tokens_limit == 456

    # When neither is passed, should be None
    limits = UsageLimits()
    assert limits.input_tokens_limit is None

    # When only new arg is set, should use it
    limits = UsageLimits(input_tokens_limit=100)
    assert limits.input_tokens_limit == 100


Review comment (suggestion):
it's weird that we need ignores... the arg types are fine, right? if we really need them use `# pyright: ignore[specificError]`

Code after:
        return 'result c'

    # Run with tool call limit of 6; expecting an error when trying to execute 3 more tools
    with pytest.raises(
        UsageLimitExceeded,
        match=re.escape('The next tool call(s) would exceed the tool_calls_limit of 6 (tool_calls=8).'),
    ):
        await agent.run('Use tools', usage_limits=UsageLimits(tool_calls_limit=6))

    # Only the first batch of 5 tools should have executed
    assert len(executed_tools) == 5


def test_usage_unknown_provider():
    assert RequestUsage.extract({}, provider='unknown', provider_url='', provider_fallback='') == RequestUsage()


def test_usage_limits_preserves_explicit_zero():
    """Test that explicit 0 token limits are preserved and not replaced by deprecated fallbacks."""
    # When input_tokens_limit=0 and deprecated request_tokens_limit is also set,
    # the explicit 0 should be preserved (not overwritten by the deprecated fallback).
    # We ignore type errors below because overloads don't allow mixing current and deprecated args.
    limits = UsageLimits(input_tokens_limit=0, request_tokens_limit=123)  # pyright: ignore[reportCallIssue]
    assert limits.input_tokens_limit == 0

    limits = UsageLimits(output_tokens_limit=0, response_tokens_limit=456)  # pyright: ignore[reportCallIssue]
    assert limits.output_tokens_limit == 0

    # When only deprecated arg is passed, should use it as fallback
    limits = UsageLimits(request_tokens_limit=123)  # pyright: ignore[reportDeprecated]
    assert limits.input_tokens_limit == 123

    limits = UsageLimits(response_tokens_limit=456)  # pyright: ignore[reportDeprecated]
    assert limits.output_tokens_limit == 456

    # When neither is passed, should be None
    limits = UsageLimits()
    assert limits.input_tokens_limit is None

    # When only current arg is set, should use it
    limits = UsageLimits(input_tokens_limit=100)
    assert limits.input_tokens_limit == 100