diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py index 37a0daa1d5..983ad73980 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py @@ -21,43 +21,51 @@ class TestBedrockFilesIntegration: file_id = "s3://test-bucket/test-file.jsonl" expected_content = b'{"recordId": "request-1", "modelInput": {}, "modelOutput": {}}' - # Mock the bedrock_files_instance.file_content method - with patch( - "litellm.files.main.bedrock_files_instance.file_content", - new_callable=AsyncMock, - ) as mock_file_content: - # Create a mock HttpxBinaryResponseContent response - import httpx + # Mock AWS credentials + with patch.dict( + "os.environ", + { + "AWS_ACCESS_KEY_ID": "test-access-key", + "AWS_SECRET_ACCESS_KEY": "test-secret-key", + }, + ): + # Mock the bedrock_files_instance.file_content method + with patch( + "litellm.files.main.bedrock_files_instance.file_content", + new_callable=AsyncMock, + ) as mock_file_content: + # Create a mock HttpxBinaryResponseContent response + import httpx - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request( - method="GET", url="s3://test-bucket/test-file.jsonl" - ), - ) - mock_file_content.return_value = HttpxBinaryResponseContent( - response=mock_response - ) + mock_response = httpx.Response( + status_code=200, + content=expected_content, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request( + method="GET", url="s3://test-bucket/test-file.jsonl" + ), + ) + mock_file_content.return_value = HttpxBinaryResponseContent( + response=mock_response + ) - # Call litellm.afile_content - result = await litellm.afile_content( - file_id=file_id, - custom_llm_provider="bedrock", - aws_region_name="us-west-2", - ) + # Call litellm.afile_content + result = await litellm.afile_content( + file_id=file_id, + custom_llm_provider="bedrock", + aws_region_name="us-west-2", + ) - # Verify the result - assert isinstance(result, HttpxBinaryResponseContent) - assert result.response.content == expected_content - assert result.response.status_code == 200 + # Verify the result + assert isinstance(result, HttpxBinaryResponseContent) + assert result.response.content == expected_content + assert result.response.status_code == 200 - # Verify the mock was called with correct parameters - mock_file_content.assert_called_once() - call_kwargs = mock_file_content.call_args.kwargs - assert call_kwargs["_is_async"] is True - assert call_kwargs["file_content_request"]["file_id"] == file_id + # Verify the mock was called with correct parameters + mock_file_content.assert_called_once() + call_kwargs = mock_file_content.call_args.kwargs + assert call_kwargs["_is_async"] is True + assert call_kwargs["file_content_request"]["file_id"] == file_id @pytest.mark.asyncio async def test_litellm_afile_content_bedrock_provider_with_unified_file_id(self): @@ -72,39 +80,47 @@ class TestBedrockFilesIntegration: expected_content = b'{"recordId": "request-1", "modelInput": {}, "modelOutput": {}}' - # Mock the bedrock_files_instance.file_content method - with patch( - "litellm.files.main.bedrock_files_instance.file_content", - new_callable=AsyncMock, - ) as mock_file_content: - # Create a mock HttpxBinaryResponseContent response - import httpx + # Mock AWS credentials + with patch.dict( + "os.environ", + { + "AWS_ACCESS_KEY_ID": "test-access-key", + "AWS_SECRET_ACCESS_KEY": "test-secret-key", + }, + ): + # Mock the bedrock_files_instance.file_content method + with patch( + "litellm.files.main.bedrock_files_instance.file_content", + new_callable=AsyncMock, + ) as mock_file_content: + # Create a mock HttpxBinaryResponseContent response + import httpx - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url=s3_uri), - ) - mock_file_content.return_value = HttpxBinaryResponseContent( - response=mock_response - ) + mock_response = httpx.Response( + status_code=200, + content=expected_content, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request(method="GET", url=s3_uri), + ) + mock_file_content.return_value = HttpxBinaryResponseContent( + response=mock_response + ) - # Call litellm.afile_content with unified file ID - result = await litellm.afile_content( - file_id=encoded_file_id, - custom_llm_provider="bedrock", - aws_region_name="us-west-2", - ) + # Call litellm.afile_content with unified file ID + result = await litellm.afile_content( + file_id=encoded_file_id, + custom_llm_provider="bedrock", + aws_region_name="us-west-2", + ) - # Verify the result - assert isinstance(result, HttpxBinaryResponseContent) - assert result.response.content == expected_content - assert result.response.status_code == 200 + # Verify the result + assert isinstance(result, HttpxBinaryResponseContent) + assert result.response.content == expected_content + assert result.response.status_code == 200 - # Verify the mock was called - the handler should extract S3 URI from unified file ID - mock_file_content.assert_called_once() - call_kwargs = mock_file_content.call_args.kwargs - assert call_kwargs["_is_async"] is True - # The handler extracts S3 URI from the unified file ID - assert call_kwargs["file_content_request"]["file_id"] == encoded_file_id + # Verify the mock was called - the handler should extract S3 URI from unified file ID + mock_file_content.assert_called_once() + call_kwargs = mock_file_content.call_args.kwargs + assert call_kwargs["_is_async"] is True + # The handler extracts S3 URI from the unified file ID + assert call_kwargs["file_content_request"]["file_id"] == encoded_file_id diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py index 723594dc39..529ebfd8bd 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py @@ -18,26 +18,13 @@ class TestVertexAIFilesIntegration: file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" expected_content = b"test file content" - # Mock the vertex_ai_files_instance.file_content method + # Mock the GCS download method to prevent actual GCS calls with patch( - "litellm.files.main.vertex_ai_files_instance.file_content", + "litellm.llms.vertex_ai.files.handler.VertexAIFilesHandler.download_gcs_object", new_callable=AsyncMock, - ) as mock_file_content: - # Create a mock HttpxBinaryResponseContent response - import httpx - - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request( - method="GET", url="gs://test-bucket/test-file.txt" - ), - ) - mock_file_content.return_value = HttpxBinaryResponseContent( - response=mock_response - ) - + ) as mock_download: + mock_download.return_value = expected_content + # Call litellm.afile_content result = await litellm.afile_content( file_id=file_id, @@ -52,13 +39,8 @@ class TestVertexAIFilesIntegration: assert result.response.content == expected_content assert result.response.status_code == 200 - # Verify the mock was called with correct parameters - mock_file_content.assert_called_once() - call_kwargs = mock_file_content.call_args.kwargs - assert call_kwargs["_is_async"] is True - assert call_kwargs["file_content_request"]["file_id"] == file_id - assert call_kwargs["vertex_project"] == "test-project" - assert call_kwargs["vertex_location"] == "us-central1" + # Verify the mock was called + mock_download.assert_called_once() def test_litellm_file_content_vertex_ai_provider(self): """Test litellm.file_content with vertex_ai provider (sync)""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index 573e095606..e80dbb85fb 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -98,9 +98,9 @@ class TestCreateToolFunction: assert callable(func) - with patch(GET_ASYNC_CLIENT_TARGET) as mock_client: + with patch(GET_ASYNC_CLIENT_TARGET) as mock_get_client: async_client = _create_mock_client("post", "verified") - mock_client.return_value = async_client + mock_get_client.return_value = async_client result = await func(**{"2fa-code": "123456"}) assert result == "verified" diff --git a/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py b/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py index 0607b0de98..d322e7ea7b 100644 --- a/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py @@ -49,11 +49,13 @@ def setup_and_teardown(): Standard LiteLLM fixture that reloads litellm before every function to speed up testing by removing callbacks being chained. """ - import importlib import asyncio + import sys - # Reload litellm to ensure clean state - importlib.reload(litellm) + # Only reload if litellm is already in sys.modules + if "litellm" in sys.modules: + import importlib + importlib.reload(litellm) # Set up async loop loop = asyncio.get_event_loop_policy().new_event_loop() diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 751a903387..c1f5d348e9 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -55,7 +55,7 @@ example_embedding_result = { def mock_patch_aembedding(): return mock.patch( - "litellm.proxy.proxy_server.llm_router.aembedding", + "litellm.aembedding", return_value=example_embedding_result, ) @@ -694,7 +694,7 @@ def test_embedding_input_array_of_tokens(mock_aembedding, client_no_auth): # Assert that aembedding was called, and that input was not modified mock_aembedding.assert_called_once() call_args, call_kwargs = mock_aembedding.call_args - assert call_kwargs["model"] == "vllm_embed_model" + assert call_kwargs["model"] == "hosted_vllm/embed_model" # Model name is transformed by router assert call_kwargs["input"] == [[2046, 13269, 158208]] assert response.status_code == 200