diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py index 60b8582089..31bcaff532 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py @@ -23,16 +23,16 @@ def mock_httpx_client(): with mock.patch( "litellm_enterprise.enterprise_callbacks.send_emails.resend_email.get_async_httpx_client" ) as mock_client: - # Create a mock response - mock_response = mock.AsyncMock(spec=Response) + + mock_response = mock.Mock(spec=Response) mock_response.status_code = 200 mock_response.json.return_value = {"id": "test_email_id"} + mock_response.raise_for_status.return_value = None - # Create a mock client mock_async_client = mock.AsyncMock() mock_async_client.post.return_value = mock_response - mock_client.return_value = mock_async_client + mock_client.return_value = mock_async_client yield mock_async_client @@ -86,7 +86,9 @@ async def test_send_email_missing_api_key(mock_httpx_client): html_body = "

Test email body

" # Mock the response to avoid making real HTTP requests - mock_response = mock.AsyncMock(spec=Response) + mock_response = mock.Mock(spec=Response) + mock_response.raise_for_status.return_value = None + mock_response.status_code = 200 mock_response.json.return_value = {"id": "test_email_id"} mock_httpx_client.post.return_value = mock_response @@ -118,7 +120,9 @@ async def test_send_email_multiple_recipients(mock_env_vars, mock_httpx_client): html_body = "

Test email body

" # Mock the response to avoid making real HTTP requests - mock_response = mock.AsyncMock(spec=Response) + mock_response = mock.Mock(spec=Response) + mock_response.raise_for_status.return_value = None + mock_response.status_code = 200 mock_response.json.return_value = {"id": "test_email_id"} mock_httpx_client.post.return_value = mock_response diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py index 2b0bb31751..f8b5e40bf3 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py @@ -14,8 +14,26 @@ from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import ( @pytest.fixture def mock_env_vars(): - with mock.patch.dict(os.environ, {"SENDGRID_API_KEY": "test_api_key"}): + # Store original values + original_api_key = os.environ.get("SENDGRID_API_KEY") + original_sender_email = os.environ.get("SENDGRID_SENDER_EMAIL") + + # Set test API key and remove SENDGRID_SENDER_EMAIL to ensure isolation + os.environ["SENDGRID_API_KEY"] = "test_api_key" + if "SENDGRID_SENDER_EMAIL" in os.environ: + del os.environ["SENDGRID_SENDER_EMAIL"] + + try: yield + finally: + # Restore original values + if original_api_key is not None: + os.environ["SENDGRID_API_KEY"] = original_api_key + elif "SENDGRID_API_KEY" in os.environ: + del os.environ["SENDGRID_API_KEY"] + + if original_sender_email is not None: + os.environ["SENDGRID_SENDER_EMAIL"] = original_sender_email @pytest.fixture @@ -23,14 +41,16 @@ def mock_httpx_client(): with mock.patch( "litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email.get_async_httpx_client" ) as mock_client: - mock_response = mock.AsyncMock(spec=Response) + + mock_response = mock.Mock(spec=Response) mock_response.status_code = 202 mock_response.text = "accepted" + mock_response.raise_for_status.return_value = None mock_async_client = mock.AsyncMock() mock_async_client.post.return_value = mock_response - mock_client.return_value = mock_async_client + mock_client.return_value = mock_async_client yield mock_async_client @@ -62,18 +82,12 @@ async def test_send_email_success(mock_env_vars, mock_httpx_client): @pytest.mark.asyncio -async def test_send_email_missing_api_key(mock_httpx_client): - # Remove the API key from environment before initializing logger +async def test_send_email_missing_api_key(): original_key = os.environ.pop("SENDGRID_API_KEY", None) - + try: logger = SendGridEmailLogger() - # Mock the response to avoid making real HTTP requests - mock_response = mock.AsyncMock(spec=Response) - mock_response.status_code = 401 - mock_httpx_client.post.return_value = mock_response - with pytest.raises(ValueError): await logger.send_email( from_email="test@example.com", @@ -81,10 +95,7 @@ async def test_send_email_missing_api_key(mock_httpx_client): subject="Test Subject", html_body="

Test email body

", ) - - mock_httpx_client.post.assert_not_called() finally: - # Restore the original key if it existed if original_key is not None: os.environ["SENDGRID_API_KEY"] = original_key @@ -98,10 +109,10 @@ async def test_send_email_multiple_recipients(mock_env_vars, mock_httpx_client): subject = "Test Subject" html_body = "

Test email body

" - # Mock the response to avoid making real HTTP requests - mock_response = mock.AsyncMock(spec=Response) + mock_response = mock.Mock(spec=Response) mock_response.status_code = 202 mock_response.text = "accepted" + mock_response.raise_for_status.return_value = None mock_httpx_client.post.return_value = mock_response await logger.send_email( diff --git a/tests/test_litellm/secret_managers/test_secret_managers_main.py b/tests/test_litellm/secret_managers/test_secret_managers_main.py index 9246b2a9ef..cadbd883cc 100644 --- a/tests/test_litellm/secret_managers/test_secret_managers_main.py +++ b/tests/test_litellm/secret_managers/test_secret_managers_main.py @@ -63,29 +63,31 @@ def test_oidc_google_success(mock_get_http_handler, mock_oidc_cache): @patch("litellm.secret_managers.main.oidc_cache") -def test_oidc_google_cached(mock_oidc_cache): +@patch("litellm.secret_managers.main._get_oidc_http_handler") +def test_oidc_google_cached(mock_get_http_handler, mock_oidc_cache): mock_oidc_cache.get_cache.return_value = "cached_token" secret_name = "oidc/google/[invalid url, do not cite]" - with patch("litellm.secret_managers.main._get_oidc_http_handler") as mock_get_http: - result = get_secret(secret_name) + result = get_secret(secret_name) - assert result == "cached_token", f"Expected cached token, got {result}" - mock_oidc_cache.get_cache.assert_called_with(key=secret_name) - mock_get_http.assert_not_called() + assert result == "cached_token", f"Expected cached token, got {result}" + mock_oidc_cache.get_cache.assert_called_with(key=secret_name) + # Verify HTTP handler was never called since we had a cached token + mock_get_http_handler.assert_not_called() @patch("litellm.secret_managers.main.oidc_cache") -def test_oidc_google_failure(mock_oidc_cache): +@patch("litellm.secret_managers.main._get_oidc_http_handler") +def test_oidc_google_failure(mock_get_http_handler, mock_oidc_cache): mock_handler = MockHTTPHandler(timeout=600.0) mock_handler.status_code = 400 + mock_get_http_handler.return_value = mock_handler + mock_oidc_cache.get_cache.return_value = None + + secret_name = "oidc/google/https://example.com/api" - with patch("litellm.secret_managers.main._get_oidc_http_handler", return_value=mock_handler): - mock_oidc_cache.get_cache.return_value = None - secret_name = "oidc/google/https://example.com/api" - - with pytest.raises(ValueError, match="Google OIDC provider failed"): - get_secret(secret_name) + with pytest.raises(ValueError, match="Google OIDC provider failed"): + get_secret(secret_name) def test_oidc_circleci_success(monkeypatch):