From 31787245797e303f56f1fd528ee688ef9ae60cf6 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 8 Aug 2024 11:28:19 -0700
Subject: [PATCH] add tests to make sure correct vertex ai route is used

---
 .../tests/test_amazing_vertex_completion.py   | 55 +++++++++++++++++++
 1 file changed, 55 insertions(+)

diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py
index bad2428fbe2..a7327bde4e7 100644
--- a/litellm/tests/test_amazing_vertex_completion.py
+++ b/litellm/tests/test_amazing_vertex_completion.py
@@ -1969,3 +1969,58 @@ def test_prompt_factory_nested():
         assert isinstance(
             message["parts"][0]["text"], str
         ), "'text' value not a string."
+
+
+def test_get_token_url():
+    from litellm.llms.vertex_httpx import VertexLLM
+
+    vertex_llm = VertexLLM()
+    vertex_ai_project = "adroit-crow-413218"
+    vertex_ai_location = "us-central1"
+    json_obj = get_vertex_ai_creds_json()
+    vertex_credentials = json.dumps(json_obj)
+
+    should_use_v1beta1_features = vertex_llm.is_using_v1beta1_features(
+        optional_params={"cached_content": "hi"}
+    )
+
+    assert should_use_v1beta1_features is True
+
+    _, url = vertex_llm._get_token_and_url(
+        vertex_project=vertex_ai_project,
+        vertex_location=vertex_ai_location,
+        vertex_credentials=vertex_credentials,
+        gemini_api_key="",
+        custom_llm_provider="vertex_ai_beta",
+        should_use_v1beta1_features=should_use_v1beta1_features,
+        api_base=None,
+        model="",
+        stream=False,
+    )
+
+    print("url=", url)
+
+    assert "/v1beta1/" in url
+
+    should_use_v1beta1_features = vertex_llm.is_using_v1beta1_features(
+        optional_params={"temperature": 0.1}
+    )
+
+    _, url = vertex_llm._get_token_and_url(
+        vertex_project=vertex_ai_project,
+        vertex_location=vertex_ai_location,
+        vertex_credentials=vertex_credentials,
+        gemini_api_key="",
+        custom_llm_provider="vertex_ai_beta",
+        should_use_v1beta1_features=should_use_v1beta1_features,
+        api_base=None,
+        model="",
+        stream=False,
+    )
+
+    print("url for normal request", url)
+
+    assert "v1beta1" not in url
+    assert "/v1/" in url
+
+    pass