From 1947c65081faa0732dfccece9b3fc544a1084742 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 02:27:18 -0700 Subject: [PATCH] test(aws): type the new partition test parameters --- tests/test_litellm/integrations/test_s3_v2.py | 2 +- .../litellm_core_utils/test_aws_partition.py | 18 +++++++++--------- .../bedrock/batches/test_transformation.py | 4 ++-- .../llms/bedrock/test_base_aws_llm.py | 15 +++++++++------ .../test_sagemaker_chat_transformation.py | 4 ++-- .../test_passthrough_endpoints_common_utils.py | 2 +- .../test_aws_secret_manager_v2.py | 12 ++++++++---- 7 files changed, 32 insertions(+), 25 deletions(-) diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index e4286fc43fa..51671d5101e 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -1883,5 +1883,5 @@ def _s3_logger_for_region(region_name: str) -> S3Logger: ), ], ) -def test_build_object_url_uses_partition_dns_suffix(region_name, expected_url): +def test_build_object_url_uses_partition_dns_suffix(region_name: str, expected_url: str) -> None: assert _s3_logger_for_region(region_name)._build_object_url("2025-01-01/key.json") == expected_url diff --git a/tests/test_litellm/litellm_core_utils/test_aws_partition.py b/tests/test_litellm/litellm_core_utils/test_aws_partition.py index fcc3e370073..3594d3c354c 100644 --- a/tests/test_litellm/litellm_core_utils/test_aws_partition.py +++ b/tests/test_litellm/litellm_core_utils/test_aws_partition.py @@ -43,7 +43,7 @@ from litellm.llms.sagemaker.chat.transformation import SagemakerChatConfig ("", "aws", "amazonaws.com"), ], ) -def test_partition_lookup(region, partition, dns_suffix): +def test_partition_lookup(region: str | None, partition: str, dns_suffix: str) -> None: assert get_aws_partition(region) == AwsPartition(partition=partition, dns_suffix=dns_suffix) assert get_aws_dns_suffix(region) == dns_suffix assert get_aws_arn_prefix(region) == f"arn:{partition}:" @@ -61,7 +61,7 @@ def test_partition_lookup(region, partition, dns_suffix): ("arn:aws:iam::123456789012:role/foo", False), ], ) -def test_contains_bedrock_arn(value, expected): +def test_contains_bedrock_arn(value: str, expected: bool) -> None: assert contains_bedrock_arn(value) is expected @@ -76,7 +76,7 @@ def test_contains_bedrock_arn(value, expected): ("arn:aws:iam::123456789012:role/foo", False), ], ) -def test_is_bedrock_arn(value, expected): +def test_is_bedrock_arn(value: str, expected: bool) -> None: assert is_bedrock_arn(value) is expected @@ -90,7 +90,7 @@ def test_is_bedrock_arn(value, expected): ("arnaws:bedrock", False), ], ) -def test_contains_aws_arn(value, expected): +def test_contains_aws_arn(value: str, expected: bool) -> None: assert contains_aws_arn(value) is expected @@ -157,14 +157,14 @@ ENDPOINT_BUILDERS: Final = { @pytest.fixture(autouse=True) -def _clear_aws_env(monkeypatch): +def _clear_aws_env(monkeypatch: pytest.MonkeyPatch) -> None: for env_var in ("AWS_BEDROCK_RUNTIME_ENDPOINT", "AWS_REGION", "AWS_DEFAULT_REGION", "AWS_REGION_NAME"): monkeypatch.delenv(env_var, raising=False) @pytest.mark.parametrize("region", ["cn-north-1", "cn-northwest-1"]) @pytest.mark.parametrize("builder_name", sorted(ENDPOINT_BUILDERS)) -def test_every_endpoint_builder_respects_cn_partition(builder_name, region): +def test_every_endpoint_builder_respects_cn_partition(builder_name: str, region: str) -> None: url = ENDPOINT_BUILDERS[builder_name](region) hostname = urlparse(url).hostname assert hostname is not None @@ -175,7 +175,7 @@ def test_every_endpoint_builder_respects_cn_partition(builder_name, region): @pytest.mark.parametrize("region", ["us-east-1", "us-gov-west-1"]) @pytest.mark.parametrize("builder_name", sorted(ENDPOINT_BUILDERS)) -def test_every_endpoint_builder_keeps_amazonaws_com_outside_cn(builder_name, region): +def test_every_endpoint_builder_keeps_amazonaws_com_outside_cn(builder_name: str, region: str) -> None: url = ENDPOINT_BUILDERS[builder_name](region) hostname = urlparse(url).hostname assert hostname is not None @@ -194,9 +194,9 @@ def _fstring_literal_offenders(needle: str) -> list[str]: ] -def test_no_fstring_hardcodes_the_commercial_dns_suffix(): +def test_no_fstring_hardcodes_the_commercial_dns_suffix() -> None: assert _fstring_literal_offenders("amazonaws.com") == [] -def test_no_fstring_hardcodes_the_commercial_arn_prefix(): +def test_no_fstring_hardcodes_the_commercial_arn_prefix() -> None: assert _fstring_literal_offenders("arn:aws:") == [] diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index 664cac6c145..7e5716a7495 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -787,7 +787,7 @@ class TestBedrockBatchesContract(BatchesConfigContractTests): expected_retrieve_status = "completed" -def test_get_complete_batch_url_cn_partition(config): +def test_get_complete_batch_url_cn_partition(config: BedrockBatchesConfig) -> None: url = config.get_complete_batch_url( api_base=None, api_key=None, @@ -812,7 +812,7 @@ def test_get_complete_batch_url_cn_partition(config): ), ], ) -def test_retrieve_request_accepts_partition_arns(config, arn, expected_prefix): +def test_retrieve_request_accepts_partition_arns(config: BedrockBatchesConfig, arn: str, expected_prefix: str) -> None: with patch.object(config.common_utils, "sign_aws_request") as mock_sign: mock_sign.return_value = ({"Authorization": "signed"}, b"") result = config.transform_retrieve_batch_request( diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index 81321ce73ec..7d07ac947b1 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -1436,8 +1436,11 @@ def test_build_sts_client_kwargs(env, aws_sts_endpoint, ssl_verify, expected): ], ) def test_resolve_sts_region_configured_region_fallback( - env, aws_sts_endpoint, aws_region_name, expected_region -): + env: dict[str, str], + aws_sts_endpoint: str | None, + aws_region_name: str | None, + expected_region: str | None, +) -> None: with patch.dict(os.environ, env, clear=True): assert ( BaseAWSLLM._resolve_sts_region( @@ -1448,7 +1451,7 @@ def test_resolve_sts_region_configured_region_fallback( ) -def test_build_sts_client_kwargs_configured_region_fallback(): +def test_build_sts_client_kwargs_configured_region_fallback() -> None: base_aws_llm = BaseAWSLLM() with patch.dict(os.environ, {}, clear=True): assert base_aws_llm._build_sts_client_kwargs(aws_region_name="cn-north-1") == { @@ -1462,7 +1465,7 @@ def test_build_sts_client_kwargs_configured_region_fallback(): } -def test_assume_role_sts_client_uses_configured_cn_region(): +def test_assume_role_sts_client_uses_configured_cn_region() -> None: """arn:aws-cn roles must resolve against a cn STS endpoint, not the commercial default.""" base_aws_llm = BaseAWSLLM() mock_expiry = MagicMock() @@ -1519,7 +1522,7 @@ def test_assume_role_sts_client_uses_configured_cn_region(): ("anthropic.claude-3", None), ], ) -def test_get_aws_region_from_model_arn_partition_arns(model, expected_region): +def test_get_aws_region_from_model_arn_partition_arns(model: str, expected_region: str | None) -> None: assert BaseAWSLLM()._get_aws_region_from_model_arn(model) == expected_region @@ -1535,7 +1538,7 @@ def test_get_aws_region_from_model_arn_partition_arns(model, expected_region): ("runtime", "us-gov-west-1", "https://bedrock-runtime.us-gov-west-1.amazonaws.com"), ], ) -def test_select_default_endpoint_url_partitions(endpoint_type, region, expected): +def test_select_default_endpoint_url_partitions(endpoint_type: str, region: str, expected: str) -> None: assert ( BaseAWSLLM()._select_default_endpoint_url( endpoint_type=endpoint_type, aws_region_name=region diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py index dd067aaa4fb..697f5a7ff59 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py @@ -344,7 +344,7 @@ def test_body_model_stays_the_endpoint_name_when_hf_model_name_is_unset(monkeypa ), ], ) -def test_get_complete_url_uses_partition_dns_suffix(region, stream, expected_url): +def test_get_complete_url_uses_partition_dns_suffix(region: str, stream: bool, expected_url: str) -> None: url = SagemakerChatConfig().get_complete_url( api_base=None, api_key=None, @@ -356,7 +356,7 @@ def test_get_complete_url_uses_partition_dns_suffix(region, stream, expected_url assert url == expected_url -def test_get_complete_url_sagemaker_base_url_override_wins(): +def test_get_complete_url_sagemaker_base_url_override_wins() -> None: url = SagemakerChatConfig().get_complete_url( api_base=None, api_key=None, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py index df36326d064..078bd4dd402 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py @@ -94,7 +94,7 @@ def test_encode_bedrock_runtime_modelid_arn_edge_cases(): assert result == expected -def test_encode_bedrock_runtime_modelid_arn_partition_arns(): +def test_encode_bedrock_runtime_modelid_arn_partition_arns() -> None: endpoint = "model/arn:aws-cn:bedrock:cn-north-1:123456789012:application-inference-profile/r742sbn2zckd/converse" expected = "model/arn:aws-cn:bedrock:cn-north-1:123456789012:application-inference-profile%2Fr742sbn2zckd/converse" assert CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint) == expected diff --git a/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py b/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py index 874be247ec6..7e655b70756 100644 --- a/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py +++ b/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py @@ -85,7 +85,9 @@ async def test_write_and_read_json_secret(): assert delete_resp is not None -def _prepare_request_endpoint(monkeypatch, region_name, extra_optional_params=None): +def _prepare_request_endpoint( + monkeypatch: pytest.MonkeyPatch, region_name: str, extra_optional_params: dict[str, str] | None = None +) -> str: monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) secret_manager = AWSSecretsManagerV2(aws_region_name=region_name) endpoint_url, _headers, _body = secret_manager._prepare_request( @@ -109,11 +111,13 @@ def _prepare_request_endpoint(monkeypatch, region_name, extra_optional_params=No ("us-east-1", "https://secretsmanager.us-east-1.amazonaws.com"), ], ) -def test_prepare_request_builds_partition_endpoint(monkeypatch, region_name, expected_endpoint): +def test_prepare_request_builds_partition_endpoint( + monkeypatch: pytest.MonkeyPatch, region_name: str, expected_endpoint: str +) -> None: assert _prepare_request_endpoint(monkeypatch, region_name) == expected_endpoint -def test_prepare_request_explicit_bedrock_runtime_endpoint_param_still_wins(monkeypatch): +def test_prepare_request_explicit_bedrock_runtime_endpoint_param_still_wins(monkeypatch: pytest.MonkeyPatch) -> None: endpoint_url = _prepare_request_endpoint( monkeypatch, "cn-north-1", @@ -122,7 +126,7 @@ def test_prepare_request_explicit_bedrock_runtime_endpoint_param_still_wins(monk assert endpoint_url == "https://secretsmanager.my-vpce.example.com" -def test_prepare_request_env_bedrock_runtime_endpoint_still_wins(monkeypatch): +def test_prepare_request_env_bedrock_runtime_endpoint_still_wins(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv( "AWS_BEDROCK_RUNTIME_ENDPOINT", "https://bedrock-runtime.eu-west-1.amazonaws.com" )