From b6be6ef574c9ac5f928c72d08f8b383a74ed1310 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 21 Nov 2023 18:02:33 -0800
Subject: [PATCH] (test) refactor test

---
 litellm/proxy/tests/test_q.py | 70 +++++++++++++++++++++--------------
 1 file changed, 43 insertions(+), 27 deletions(-)

diff --git a/litellm/proxy/tests/test_q.py b/litellm/proxy/tests/test_q.py
index 0fc7b4b064..37c42fb3e4 100644
--- a/litellm/proxy/tests/test_q.py
+++ b/litellm/proxy/tests/test_q.py
@@ -1,46 +1,68 @@
 import requests
 import time
 
-# Step 1 Add a config to the proxy, generate a temp key 
-config = {
+# Set the base URL as needed
+# base_url = "https://litellm-api.onrender.com"
+# Uncomment the line below if you want to switch to the local server
+base_url = "http://0.0.0.0:8000"
 
+# Step 1 Add a config to the proxy, generate a temp key
+config = {
+  "model_list": [
+    {
+      "model_name": "gpt-3.5-turbo",
+      "litellm_params": {
+        "model": "gpt-3.5-turbo",
+        "api_key": "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"
+      }
+    },
+    {
+      "model_name": "gpt-3.5-turbo",
+      "litellm_params": {
+        "model": "azure/chatgpt-v-2",
+        "api_key": "6314c6dc63f448c9873844297f408c74",
+        "api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/",
+        "api_version": "2023-07-01-preview"
+      }
+    }
+  ]
 }
 
 response = requests.post(
-    url = "http://0.0.0.0:8000/key/generate",
+    url=f"{base_url}/key/generate",
     json={
         "config": config,
-        "duration": "30d" # default to 30d, set it to 30m if you want a temp key
+        "duration": "30d"  # default to 30d, set it to 30m if you want a temp key
     },
     headers={
-        "Authorization": "Bearer sk-hosted-litellm"
+        "Authorization": "Bearer sk-1234"
     }
 )
 
-print("\nresponse from generating key", response.json())
+print("\nresponse from generating key", response.text)
 
 generated_key = response.json()["key"]
 print("\ngenerated key for proxy", generated_key)
 
-
 # Step 2: Queue a request to the proxy, using your generated_key
 job_response = requests.post(
-    url = "http://0.0.0.0:8000/queue/request",
+    url=f"{base_url}/queue/request",
     json={
-            'model': 'gpt-3.5-turbo',
-            'messages': [
-                {'role': 'system', 'content': f'You are a helpful assistant. What is your name'},
-            ],
+        'model': 'gpt-3.5-turbo',
+        'messages': [
+            {'role': 'system', 'content': f'You are a helpful assistant. What is your name'},
+        ],
     },
     headers={
         "Authorization": f"Bearer {generated_key}"
     }
 )
 
+print("\nResponse from creating job", job_response.text)
 job_response = job_response.json()
-job_id  = job_response["id"]
+job_id = job_response["id"]
 polling_url = job_response["url"]
-polling_url = f"http://0.0.0.0:8000{polling_url}"
+polling_url = f"{base_url}{polling_url}"
 print("\nCreated Job, Polling Url", polling_url)
 
 # Step 3: Poll the request
@@ -48,14 +70,14 @@ while True:
     try:
         print("\nPolling URL", polling_url)
         polling_response = requests.get(
-                url=polling_url, 
-                headers={
-                    "Authorization": f"Bearer {generated_key}"
-                }
-            )
+            url=polling_url,
+            headers={
+                "Authorization": f"Bearer {generated_key}"
+            }
+        )
+        print("\nResponse from polling url", polling_response.text)
         polling_response = polling_response.json()
-        print("\nResponse from polling url", polling_response)
-        status = polling_response["status"]
+        status = polling_response.get("status", None)
         if status == "finished":
             llm_response = polling_response["result"]
             print("LLM Response")
@@ -65,9 +87,3 @@ while True:
     except Exception as e:
         print("got exception in polling", e)
         break
-
-
-
-
-
-