diff --git a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py b/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py deleted file mode 100644 index 0b11a66c100..00000000000 --- a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py +++ /dev/null @@ -1,158 +0,0 @@ -import json -import os -from unittest.mock import Mock, patch -import pytest - - -import litellm -from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler - -# Mock response for Bedrock image generation -mock_image_response = {"images": ["base64_encoded_image_data"], "error": None} - - -class TestBedrockImageGeneration: - def test_image_generation_with_api_key_bearer_token(self): - """Test image generation with bearer token authentication""" - test_api_key = "test-bearer-token-12345" - model = "bedrock/stability.sd3-large-v1:0" - prompt = "A cute baby sea otter" - - with patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.image_generation" - ) as mock_bedrock_image_gen: - # Setup mock response - mock_image_response_obj = litellm.ImageResponse() - mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}] - mock_bedrock_image_gen.return_value = mock_image_response_obj - - response = litellm.image_generation( - model=model, - prompt=prompt, - aws_region_name="us-west-2", - api_key=test_api_key, - ) - - assert response is not None - assert len(response.data) > 0 - - mock_bedrock_image_gen.assert_called_once() - for call in mock_bedrock_image_gen.call_args_list: - if "headers" in call.kwargs: - headers = call.kwargs["headers"] - if ( - "Authorization" in headers - and headers["Authorization"] == f"Bearer {test_api_key}" - ): - break - - def test_image_generation_with_env_variable_bearer_token(self, monkeypatch): - """Test image generation with bearer token from environment variable""" - test_api_key = "env-bearer-token-12345" - model = "bedrock/stability.sd3-large-v1:0" - prompt = "A cute baby sea otter" - - # Mock the environment variable - with ( - patch.dict(os.environ, {"AWS_BEARER_TOKEN_BEDROCK": test_api_key}), - patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.image_generation" - ) as mock_bedrock_image_gen, - ): - - mock_image_response_obj = litellm.ImageResponse() - mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}] - mock_bedrock_image_gen.return_value = mock_image_response_obj - - response = litellm.image_generation( - model=model, prompt=prompt, aws_region_name="us-west-2" - ) - - assert response is not None - assert len(response.data) > 0 - - mock_bedrock_image_gen.assert_called_once() - for call in mock_bedrock_image_gen.call_args_list: - if "headers" in call.kwargs: - headers = call.kwargs["headers"] - if ( - "Authorization" in headers - and headers["Authorization"] == f"Bearer {test_api_key}" - ): - break - - @pytest.mark.asyncio - async def test_async_image_generation_with_bearer_token(self): - """Test async image generation with bearer token authentication""" - test_api_key = "async-bearer-token-12345" - model = "bedrock/stability.sd3-large-v1:0" - prompt = "A cute baby sea otter" - - with patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.async_image_generation" - ) as mock_async_bedrock_image_gen: - mock_image_response_obj = litellm.ImageResponse() - mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}] - mock_async_bedrock_image_gen.return_value = mock_image_response_obj - - # Call async image generation with api_key parameter - response = await litellm.aimage_generation( - model=model, - prompt=prompt, - aws_region_name="us-west-2", - api_key=test_api_key, - ) - - assert response is not None - assert len(response.data) > 0 - - mock_async_bedrock_image_gen.assert_called_once() - for call in mock_async_bedrock_image_gen.call_args_list: - if "headers" in call.kwargs: - headers = call.kwargs["headers"] - if ( - "Authorization" in headers - and headers["Authorization"] == f"Bearer {test_api_key}" - ): - break - - def test_image_generation_with_sigv4(self): - """Test image generation falls back to SigV4 auth when no bearer token is provided""" - model = "bedrock/stability.sd3-large-v1:0" - prompt = "A cute baby sea otter" - - with patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.image_generation" - ) as mock_bedrock_image_gen: - mock_image_response_obj = litellm.ImageResponse() - mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}] - mock_bedrock_image_gen.return_value = mock_image_response_obj - - response = litellm.image_generation( - model=model, prompt=prompt, aws_region_name="us-west-2" - ) - - assert response is not None - assert len(response.data) > 0 - mock_bedrock_image_gen.assert_called_once() - - -def test_image_generation_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): - """The deployment's AWS profile does not exist, so resolving SigV4 credentials - raises; a bearer-token deployment must still sign the request with the - bearer token alone.""" - from litellm.llms.bedrock.image_generation.image_handler import BedrockImageGeneration - - monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") - - request = BedrockImageGeneration()._prepare_request( - model="amazon.nova-canvas-v1:0", - prompt="A cute baby sea otter", - optional_params={"aws_region_name": "us-west-2", "aws_profile_name": "litellm-no-such-aws-profile"}, - api_base=None, - extra_headers=None, - api_key=None, - logging_obj=Mock(), - ) - - assert request.prepped.headers["Authorization"] == "Bearer env-bearer-token-12345" diff --git a/tests/test_litellm/llms/bedrock/image/test_amazon_nova_canvas_transformation.py b/tests/unit/llms/bedrock/image/test_amazon_nova_canvas_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/image/test_amazon_nova_canvas_transformation.py rename to tests/unit/llms/bedrock/image/test_amazon_nova_canvas_transformation.py diff --git a/tests/test_litellm/llms/bedrock/image/test_amazon_stability3_transformation.py b/tests/unit/llms/bedrock/image/test_amazon_stability3_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/image/test_amazon_stability3_transformation.py rename to tests/unit/llms/bedrock/image/test_amazon_stability3_transformation.py diff --git a/tests/unit/llms/bedrock/image/test_bedrock_image_bearer_token.py b/tests/unit/llms/bedrock/image/test_bedrock_image_bearer_token.py new file mode 100644 index 00000000000..599507da03d --- /dev/null +++ b/tests/unit/llms/bedrock/image/test_bedrock_image_bearer_token.py @@ -0,0 +1,21 @@ +from unittest.mock import Mock + +def test_image_generation_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): + """The deployment's AWS profile does not exist, so resolving SigV4 credentials + raises; a bearer-token deployment must still sign the request with the + bearer token alone.""" + from litellm.llms.bedrock.image_generation.image_handler import BedrockImageGeneration + + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") + + request = BedrockImageGeneration()._prepare_request( + model="amazon.nova-canvas-v1:0", + prompt="A cute baby sea otter", + optional_params={"aws_region_name": "us-west-2", "aws_profile_name": "litellm-no-such-aws-profile"}, + api_base=None, + extra_headers=None, + api_key=None, + logging_obj=Mock(), + ) + + assert request.prepped.headers["Authorization"] == "Bearer env-bearer-token-12345" diff --git a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_prepare_request.py b/tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py similarity index 100% rename from tests/test_litellm/llms/bedrock/image/test_bedrock_image_prepare_request.py rename to tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py diff --git a/tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py b/tests/unit/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py similarity index 100% rename from tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py rename to tests/unit/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py diff --git a/tests/test_litellm/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py b/tests/unit/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py rename to tests/unit/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py diff --git a/tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py similarity index 100% rename from tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/test_handler.py rename to tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py diff --git a/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py similarity index 99% rename from tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py rename to tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index f2a9af11af7..854ef92fa4b 100644 --- a/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -367,8 +367,6 @@ def test_bedrock_passthrough_region_extraction_from_inference_profile_arn(): assert ( "us-west-2" in api_base ), f"Expected region 'us-west-2' from ARN in base URL, but got: {api_base}" - - def test_bedrock_passthrough_model_id_arn_encoding(): """ Test that model_id ARNs are properly URL-encoded when used in endpoints. diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py similarity index 100% rename from tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py rename to tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py rename to tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py similarity index 100% rename from tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py rename to tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py rename to tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py diff --git a/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py b/tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py rename to tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py diff --git a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py b/tests/unit/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py rename to tests/unit/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py diff --git a/tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py b/tests/unit/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py rename to tests/unit/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py diff --git a/tests/test_litellm/llms/black_forest_labs/test_bfl_common_utils.py b/tests/unit/llms/black_forest_labs/test_bfl_common_utils.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/test_bfl_common_utils.py rename to tests/unit/llms/black_forest_labs/test_bfl_common_utils.py diff --git a/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py b/tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py rename to tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py diff --git a/tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py b/tests/unit/llms/cerebras/test_cerebras_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py rename to tests/unit/llms/cerebras/test_cerebras_chat_transformation.py diff --git a/tests/test_litellm/llms/chat/test_converse_handler.py b/tests/unit/llms/chat/test_converse_handler.py similarity index 100% rename from tests/test_litellm/llms/chat/test_converse_handler.py rename to tests/unit/llms/chat/test_converse_handler.py diff --git a/tests/test_litellm/llms/chatgpt/test_chatgpt_authenticator.py b/tests/unit/llms/chatgpt/test_chatgpt_authenticator.py similarity index 100% rename from tests/test_litellm/llms/chatgpt/test_chatgpt_authenticator.py rename to tests/unit/llms/chatgpt/test_chatgpt_authenticator.py