From b77492d574b47a0110f30209cbc43b072407c5fa Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 13 Oct 2023 21:09:08 -0700
Subject: [PATCH] (tests) add azure/embedding search

---
 litellm/tests/test_embedding.py | 45 +++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py
index e27f930bed3..4ca2a9eea54 100644
--- a/litellm/tests/test_embedding.py
+++ b/litellm/tests/test_embedding.py
@@ -22,6 +22,51 @@ def test_openai_embedding():
         pytest.fail(f"Error occurred: {e}")
 # test_openai_embedding()
 
+def test_openai_azure_embedding_simple():
+    try:
+
+        response = embedding(
+            model="azure/azure-embedding-model",
+            input=["good morning from litellm"],
+        )
+        print(response)
+
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+
+# test_openai_azure_embedding_simple()
+
+def test_openai_azure_embedding():
+    try:
+        api_key = os.environ['AZURE_API_KEY']
+        api_base = os.environ['AZURE_API_BASE']
+        api_version = os.environ['AZURE_API_VERSION']
+
+        os.environ['AZURE_API_VERSION'] = ""
+        os.environ['AZURE_API_BASE'] = ""
+        os.environ['AZURE_API_KEY'] = ""
+
+        response = embedding(
+            model="azure/azure-embedding-model",
+            input=["good morning from litellm", "this is another item"],
+            api_key=api_key,
+            api_base=api_base,
+            api_version=api_version,
+        )
+        print(response)
+
+
+        os.environ['AZURE_API_VERSION'] = api_key
+        os.environ['AZURE_API_BASE'] = api_base
+        os.environ['AZURE_API_KEY'] = api_base
+
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+
+# test_openai_azure_embedding()
+
+# test_openai_embedding()
+
 def test_cohere_embedding():
     try:
         response = embedding(