diff --git a/tests/litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio.py b/tests/litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio.py index 5e870b5af36..9fc9be7e0f3 100644 --- a/tests/litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio.py +++ b/tests/litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio.py @@ -31,3 +31,41 @@ def test_top_logprobs(): ) assert v["responseLogprobs"] is non_default_params["logprobs"] assert v["logprobs"] is non_default_params["top_logprobs"] + + +def test_get_model_for_vertex_ai_url(): + # Test case 1: Regular model name + model = "gemini-pro" + result = VertexGeminiConfig.get_model_for_vertex_ai_url(model) + assert result == "gemini-pro" + + # Test case 2: Gemini spec model with UUID + model = "gemini/ft-uuid-123" + result = VertexGeminiConfig.get_model_for_vertex_ai_url(model) + assert result == "ft-uuid-123" + + +def test_is_model_gemini_gemini_spec_model(): + # Test case 1: None input + assert VertexGeminiConfig._is_model_gemini_gemini_spec_model(None) == False + + # Test case 2: Regular model name + assert VertexGeminiConfig._is_model_gemini_gemini_spec_model("gemini-pro") == False + + # Test case 3: Gemini spec model + assert ( + VertexGeminiConfig._is_model_gemini_gemini_spec_model("gemini/custom-model") + == True + ) + + +def test_get_model_name_from_gemini_spec_model(): + # Test case 1: Regular model name + model = "gemini-pro" + result = VertexGeminiConfig._get_model_name_from_gemini_spec_model(model) + assert result == "gemini-pro" + + # Test case 2: Gemini spec model + model = "gemini/ft-uuid-123" + result = VertexGeminiConfig._get_model_name_from_gemini_spec_model(model) + assert result == "ft-uuid-123"