From 7350453aa1ce75549fdfaa9f7eb1182796cf4f66 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 20 Feb 2025 13:55:56 -0800 Subject: [PATCH] test: add unit testing to ensure bedrock region name inferred from arn on rerank --- tests/litellm/rerank_api/test_main.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/litellm/rerank_api/test_main.py b/tests/litellm/rerank_api/test_main.py index f55e05ea0f6..d3051b81f92 100644 --- a/tests/litellm/rerank_api/test_main.py +++ b/tests/litellm/rerank_api/test_main.py @@ -14,8 +14,10 @@ from litellm import rerank from litellm.llms.custom_httpx.http_handler import HTTPHandler -def test_rerank_infer_region_from_model_arn(): +def test_rerank_infer_region_from_model_arn(monkeypatch): mock_response = MagicMock() + + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") args = { "model": "bedrock/arn:aws:bedrock:us-west-2::foundation-model/amazon.rerank-v1:0", "query": "hello", @@ -45,3 +47,5 @@ def test_rerank_infer_region_from_model_arn(): ) mock_post.assert_called_once() print(f"mock_post.call_args: {mock_post.call_args.kwargs}") + assert "us-west-2" in mock_post.call_args.kwargs["url"] + assert "us-east-1" not in mock_post.call_args.kwargs["url"]