diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index 7b3c0b27842..1505c39d4a1 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -29,6 +29,42 @@ from litellm.llms.custom_httpx.async_client_cleanup import ( from litellm.proxy.db import tool_registry_writer as tool_registry_writer_module +def _reset_module_level_aws_auth_caches(): + """ + Clear module-level AWS auth state that can survive between tests. + + Bedrock/SageMaker handlers are instantiated once at import time and cache + resolved credentials on the handler instance. If a previous test resolves an + invalid or different auth flow, later tests can reuse that cached state and + bypass their local monkeypatched env setup. + """ + for module_name in ( + "litellm.main", + "litellm.files.main", + "litellm.rerank_api.main", + "litellm.realtime_api.main", + ): + try: + module = importlib.import_module(module_name) + except Exception: + continue + for attr_name in dir(module): + obj = getattr(module, attr_name) + iam_cache = getattr(obj, "iam_cache", None) + if iam_cache is None: + continue + flush_cache = getattr(iam_cache, "flush_cache", None) + if callable(flush_cache): + flush_cache() + + try: + import boto3 + + boto3.DEFAULT_SESSION = None + except Exception: + pass + + @pytest.fixture(scope="session") def isolated_aws_credentials_dir(tmp_path_factory): aws_dir = tmp_path_factory.mktemp("aws-config") @@ -54,6 +90,12 @@ def isolate_host_aws_config(monkeypatch, isolated_aws_credentials_dir): monkeypatch.delenv("AWS_DEFAULT_PROFILE", raising=False) monkeypatch.delenv("AWS_CONTAINER_CREDENTIALS_FULL_URI", raising=False) monkeypatch.delenv("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", raising=False) + monkeypatch.delenv("AWS_SESSION_TOKEN", raising=False) + monkeypatch.delenv("AWS_ROLE_ARN", raising=False) + monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("AWS_REGION_NAME", raising=False) + monkeypatch.delenv("AWS_DEFAULT_REGION", raising=False) def _run_coroutine_if_needed(result): @@ -196,6 +238,7 @@ def isolate_litellm_state(): if hasattr(litellm, "in_memory_llm_clients_cache"): litellm.in_memory_llm_clients_cache.flush_cache() image_handling_module.in_memory_cache.flush_cache() + _reset_module_level_aws_auth_caches() # Clear all callback lists to prevent cross-test contamination if hasattr(litellm, 'callbacks'): @@ -228,6 +271,7 @@ def isolate_litellm_state(): if hasattr(litellm, "in_memory_llm_clients_cache"): litellm.in_memory_llm_clients_cache.flush_cache() image_handling_module.in_memory_cache.flush_cache() + _reset_module_level_aws_auth_caches() current_module_level_client = litellm.__dict__.get("module_level_client") current_module_level_aclient = litellm.__dict__.get("module_level_aclient") diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index c241813d08a..96e998fd33c 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -18,6 +18,10 @@ import litellm from litellm import main as litellm_main +async def _async_fake_bedrock_image_details(image_url): + return "ZmFrZS1pbWFnZQ==", "image/png" + + @pytest.fixture(autouse=True) def clear_client_cache(): """ @@ -166,12 +170,31 @@ def test_completion_missing_role(openai_api_response): async def test_url_with_format_param(model, sync_mode, monkeypatch): from litellm import acompletion, completion from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.litellm_core_utils.prompt_templates import factory as prompt_factory if sync_mode: client = HTTPHandler() else: client = AsyncHTTPHandler() + # This test is about request shaping, not live image downloads. Stub the + # URL->image conversion helpers so suite-level network/client state from + # earlier tests cannot prevent the mocked provider client from being hit. + fake_base64_image = "data:image/png;base64,ZmFrZS1pbWFnZQ==" + monkeypatch.setattr( + prompt_factory, "convert_url_to_base64", lambda url: fake_base64_image + ) + monkeypatch.setattr( + prompt_factory.BedrockImageProcessor, + "get_image_details", + staticmethod(lambda image_url: ("ZmFrZS1pbWFnZQ==", "image/png")), + ) + monkeypatch.setattr( + prompt_factory.BedrockImageProcessor, + "get_image_details_async", + staticmethod(_async_fake_bedrock_image_details), + ) + args = { "model": model, "messages": [