From 41a5daf20c771559ab6f2f42f47549c94ad456c8 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 29 Aug 2024 15:50:54 -0700
Subject: [PATCH] add test for vertex sdk foward headers

---
 litellm/proxy/proxy_config.yaml                    |  8 --------
 .../proxy/tests/test_vertex_sdk_forward_headers.py | 14 ++++++++++++++
 2 files changed, 14 insertions(+), 8 deletions(-)
 create mode 100644 litellm/proxy/tests/test_vertex_sdk_forward_headers.py

diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index ac17a1d84c4..21beb965cc0 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/litellm/proxy/proxy_config.yaml
@@ -13,14 +13,6 @@ model_list:
       model: cohere/rerank-english-v3.0
       api_key: os.environ/COHERE_API_KEY
 
-general_settings:
-  enable_oauth2_proxy_auth: True
-  oauth2_config_mappings:
-    token: X-Auth-Token
-    user_id: X-Auth-Client-ID
-    team_id: X-Auth-Team-ID
-    max_budget: X-Auth-Max-Budget
-    models: X-Auth-Allowed-Models
 
 # default off mode
 litellm_settings:
diff --git a/litellm/proxy/tests/test_vertex_sdk_forward_headers.py b/litellm/proxy/tests/test_vertex_sdk_forward_headers.py
new file mode 100644
index 00000000000..b291be438ca
--- /dev/null
+++ b/litellm/proxy/tests/test_vertex_sdk_forward_headers.py
@@ -0,0 +1,14 @@
+import vertexai
+from vertexai.preview.generative_models import GenerativeModel
+
+LITE_LLM_ENDPOINT = "http://localhost:4000"
+
+vertexai.init(
+    project="adroit-crow-413218",
+    location="us-central1",
+    api_endpoint=f"{LITE_LLM_ENDPOINT}/vertex-ai",
+    api_transport="rest",
+)
+
+model = GenerativeModel(model_name="gemini-1.0-pro")
+model.generate_content("hi")