test: cleanup testing

litellm/tests/test_completion.py

@@ -2611,18 +2611,37 @@ def test_completion_azure_ad_token():
     # If you want to remove it, speak to Ishaan!
     # Ishaan will be very disappointed if this test is removed -> this is a standard way to pass api_key + the router + proxy use this
     from httpx import Client
-    from openai import AzureOpenAI
 
     from litellm import completion
-    from litellm.llms.custom_httpx.httpx_handler import HTTPHandler
 
-    response = completion(
-        model="azure/chatgpt-v-2",
-        messages=messages,
-        # api_key="my-fake-ad-token",
-        azure_ad_token=os.getenv("AZURE_API_KEY"),
-    )
-    print(response)
+    litellm.set_verbose = True
+
+    old_key = os.environ["AZURE_API_KEY"]
+    os.environ.pop("AZURE_API_KEY", None)
+
+    http_client = Client()
+
+    with patch.object(http_client, "send", new=MagicMock()) as mock_client:
+        litellm.client_session = http_client
+        try:
+            response = completion(
+                model="azure/chatgpt-v-2",
+                messages=messages,
+                azure_ad_token="my-special-token",
+            )
+            print(response)
+        except Exception as e:
+            pass
+        finally:
+            os.environ["AZURE_API_KEY"] = old_key
+
+        mock_client.assert_called_once()
+        request = mock_client.call_args[0][0]
+        print(request.method)  # This will print 'POST'
+        print(request.url)  # This will print the full URL
+        print(request.headers)  # This will print the full URL
+        auth_header = request.headers.get("Authorization")
+        assert auth_header == "Bearer my-special-token"
 
 
 def test_completion_azure_key_completion_arg():

litellm/tests/test_embedding.py

@@ -206,6 +206,9 @@ def test_openai_azure_embedding_with_oidc_and_cf():
     os.environ["AZURE_TENANT_ID"] = "17c0a27a-1246-4aa1-a3b6-d294e80e783c"
     os.environ["AZURE_CLIENT_ID"] = "4faf5422-b2bd-45e8-a6d7-46543a38acd0"
 
+    old_key = os.environ["AZURE_API_KEY"]
+    os.environ.pop("AZURE_API_KEY", None)
+
     try:
         response = embedding(
             model="azure/text-embedding-ada-002",
@@ -218,6 +221,8 @@ def test_openai_azure_embedding_with_oidc_and_cf():
 
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
+    finally:
+        os.environ["AZURE_API_KEY"] = old_key
 
 
 def test_openai_azure_embedding_optional_arg(mocker):
@@ -673,17 +678,3 @@ async def test_databricks_embeddings(sync_mode):
 #     print(response)
 
 # local_proxy_embeddings()
-
-
-def test_embedding_azure_ad_token():
-    # this tests if we can pass api_key to completion, when it's not in the env.
-    # DO NOT REMOVE THIS TEST. No MATTER WHAT Happens!
-    # If you want to remove it, speak to Ishaan!
-    # Ishaan will be very disappointed if this test is removed -> this is a standard way to pass api_key + the router + proxy use this
-
-    response = embedding(
-        model="azure/azure-embedding-model",
-        input=["good morning from litellm"],
-        azure_ad_token=os.getenv("AZURE_API_KEY"),
-    )
-    print(response)