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

Code before:
    assert cohere_profile is None

    deepseek_profile = provider.model_profile('deepseek.deepseek-r1')
    deepseek_model_profile_mock.assert_called_with('deepseek-r1')
    assert deepseek_profile is not None
    assert deepseek_profile.ignore_streamed_leading_whitespace is True

    amazon_profile = provider.model_profile('us.amazon.nova-pro-v1:0')
    amazon_model_profile_mock.assert_called_with('nova-pro')
    assert isinstance(amazon_profile, BedrockModelProfile)
    assert amazon_profile.json_schema_transformer == InlineDefsJsonSchemaTransformer
    assert amazon_profile.bedrock_supports_tool_choice is True

    amazon_profile = provider.model_profile('us.amazon.titan-text-express-v1:0')
    amazon_model_profile_mock.assert_called_with('titan-text-express')
    assert amazon_profile is not None
    assert amazon_profile.json_schema_transformer == InlineDefsJsonSchemaTransformer

    unknown_model = provider.model_profile('unknown-model')
    assert unknown_model is None

    unknown_model = provider.model_profile('unknown.unknown-model')
    assert unknown_model is None


@pytest.mark.parametrize('prefix', BEDROCK_GEO_PREFIXES)
def test_bedrock_provider_model_profile_all_geo_prefixes(env: TestEnv, prefix: str):
    """Test that all cross-region inference geo prefixes are correctly handled."""
    env.set('AWS_DEFAULT_REGION', 'us-east-1')
    provider = BedrockProvider()

    model_name = f'{prefix}anthropic.claude-sonnet-4-5-20250929-v1:0'
    profile = provider.model_profile(model_name)

    assert profile is not None, f'model_profile returned None for {model_name}'


def test_latest_bedrock_model_names_geo_prefixes_are_supported():
    """Ensure all geo prefixes used in LatestBedrockModelNames are in BEDROCK_GEO_PREFIXES.

    This test prevents adding new model names with geo prefixes that aren't handled
    by the provider's model_profile method.
    """
    model_names = get_args(LatestBedrockModelNames)

    missing_prefixes: set[str] = set()

    for model_name in model_names:
        # Model names with geo prefixes have 3+ dot-separated parts:
        # - No prefix: "anthropic.claude-xxx" (2 parts)
        # - With prefix: "us.anthropic.claude-xxx" (3 parts)

Review comment (suggestion):
tests look good now, thanks!

can we add another test like this (non-parametrized though) to show what happens when a user passes a model with a prefix that's not in the list? you can make one up like `narnia.`

Code after:
    assert cohere_profile is None

    deepseek_profile = provider.model_profile('deepseek.deepseek-r1')
    deepseek_model_profile_mock.assert_called_with('deepseek-r1')
    assert deepseek_profile is not None
    assert deepseek_profile.ignore_streamed_leading_whitespace is True

    amazon_profile = provider.model_profile('us.amazon.nova-pro-v1:0')
    amazon_model_profile_mock.assert_called_with('nova-pro')
    assert isinstance(amazon_profile, BedrockModelProfile)
    assert amazon_profile.json_schema_transformer == InlineDefsJsonSchemaTransformer
    assert amazon_profile.bedrock_supports_tool_choice is True

    amazon_profile = provider.model_profile('us.amazon.titan-text-express-v1:0')
    amazon_model_profile_mock.assert_called_with('titan-text-express')
    assert amazon_profile is not None
    assert amazon_profile.json_schema_transformer == InlineDefsJsonSchemaTransformer

    unknown_model = provider.model_profile('unknown-model')
    assert unknown_model is None

    unknown_model = provider.model_profile('unknown.unknown-model')
    assert unknown_model is None


@pytest.mark.parametrize('prefix', BEDROCK_GEO_PREFIXES)
def test_bedrock_provider_model_profile_all_geo_prefixes(env: TestEnv, prefix: str):
    """Test that all cross-region inference geo prefixes are correctly handled."""
    env.set('AWS_DEFAULT_REGION', 'us-east-1')
    provider = BedrockProvider()

    model_name = f'{prefix}.anthropic.claude-sonnet-4-5-20250929-v1:0'
    profile = provider.model_profile(model_name)

    assert profile is not None, f'model_profile returned None for {model_name}'


def test_bedrock_provider_model_profile_with_unknown_geo_prefix(env: TestEnv):
    env.set('AWS_DEFAULT_REGION', 'us-east-1')
    provider = BedrockProvider()

    model_name = 'narnia.anthropic.claude-sonnet-4-5-20250929-v1:0'
    profile = provider.model_profile(model_name)
    assert profile is None, f'model_profile returned {profile} for {model_name}'


def test_latest_bedrock_model_names_geo_prefixes_are_supported():
    """Ensure all geo prefixes used in LatestBedrockModelNames are in BEDROCK_GEO_PREFIXES.

    This test prevents adding new model names with geo prefixes that aren't handled
    by the provider's model_profile method.