(test) Router: Test Azure acompletion, stream

litellm/tests/test_router.py

@@ -139,6 +139,89 @@ def test_reading_key_from_model_list():
 # test_reading_key_from_model_list()
 
 
+
+def test_router_azure_acompletion():
+	# [PROD TEST CASE]
+	# This is 90% of the router use case, makes an acompletion call, acompletion + stream call and verifies it got a response
+	# DO NOT REMOVE THIS TEST. It's an IMP ONE. Speak to Ishaan, if you are tring to remove this
+	litellm.set_verbose=False
+	import openai
+	try:
+		print("Router Test Azure - Acompletion, Acompletion with stream")
+
+		# remove api key from env to repro how proxy passes key to router 
+		old_api_key = os.environ["AZURE_API_KEY"]
+		os.environ.pop("AZURE_API_KEY", None)
+	
+		model_list = [
+			{ 
+				"model_name": "gpt-3.5-turbo", # openai model name 
+				"litellm_params": { # params for litellm completion/embedding call 
+					"model": "azure/chatgpt-v-2", 
+					"api_key": old_api_key,
+					"api_version": os.getenv("AZURE_API_VERSION"),
+					"api_base": os.getenv("AZURE_API_BASE")
+				},
+				"tpm": 240000,
+				"rpm": 1800
+			},
+			{
+				"model_name": "gpt-3.5-turbo", # openai model name 
+				"litellm_params": { # params for litellm completion/embedding call 
+					"model": "azure/chatgpt-v-2", 
+					"api_key": old_api_key,
+					"api_version": os.getenv("AZURE_API_VERSION"),
+					"api_base": os.getenv("AZURE_API_BASE")
+				},
+				"tpm": 240000,
+				"rpm": 1800
+			}
+		]
+
+		router = Router(model_list=model_list, 
+					routing_strategy="simple-shuffle",
+					set_verbose=True
+				) # type: ignore
+		
+		async def test1():
+
+			response = await router.acompletion(
+				model="gpt-3.5-turbo",
+				messages=[
+					{
+						"role": "user",
+						"content": "hello this request will fail"
+					}
+				]
+			)
+			print("\n response", response)
+		asyncio.run(test1())
+
+		print("\n Testing streaming response")
+		async def test2():
+			response = await router.acompletion(
+				model="gpt-3.5-turbo",
+				messages=[
+					{
+						"role": "user",
+						"content": "hello this request will fail"
+					}
+				],
+				stream=True
+			)
+			async for chunk in response:
+				if chunk is not None:
+					print(chunk)
+		asyncio.run(test2())
+		print("\n Passed Streaming")
+		os.environ["AZURE_API_KEY"] = old_api_key
+		router.reset()
+	except Exception as e:
+		os.environ["AZURE_API_KEY"] = old_api_key
+		print(f"FAILED TEST")
+		pytest.fail(f"Got unexpected exception on router! - {e}")
+test_router_azure_acompletion()
+
 ### FUNCTION CALLING 
 
 def test_function_calling():