fix: add more tests

Author: Abraham Lazaro Martinez

py/plugins/vertex-ai/tests/model_garden/test_client.py

@@ -24,23 +24,18 @@
 @patch('google.auth.default')
 @patch('google.auth.transport.requests.Request')
 @patch('openai.OpenAI')
-def test_client_initialization_with_explicit_project_id(
-    mock_openai_cls, mock_request_cls, mock_default_auth
-):
+def test_client_initialization_with_explicit_project_id(mock_openai_cls, mock_request_cls, mock_default_auth):
     """Unittests for init client."""
-    mock_location = "location"
-    mock_project_id = "project_id"
-    mock_token = "token"
+    mock_location = 'location'
+    mock_project_id = 'project_id'
+    mock_token = 'token'
 
     mock_credentials = MagicMock()
     mock_credentials.token = mock_token
 
-    mock_default_auth.return_value = (mock_credentials, "project_id")
+    mock_default_auth.return_value = (mock_credentials, 'project_id')
 
-    client_instance = OpenAIClient(
-        location=mock_location,
-        project_id=mock_project_id
-    )
+    client_instance = OpenAIClient(location=mock_location, project_id=mock_project_id)
 
     mock_default_auth.assert_called_once()
     mock_credentials.refresh.assert_called_once()
@@ -52,17 +47,15 @@ def test_client_initialization_with_explicit_project_id(
 @patch('google.auth.default')
 @patch('google.auth.transport.requests.Request')
 @patch('openai.OpenAI')
-def test_client_initialization_without_explicit_project_id(
-    mock_openai_cls, mock_request_cls, mock_default_auth
-):
+def test_client_initialization_without_explicit_project_id(mock_openai_cls, mock_request_cls, mock_default_auth):
     """Unittests for init client."""
-    mock_location = "location"
-    mock_token = "token"
+    mock_location = 'location'
+    mock_token = 'token'
 
     mock_credentials = MagicMock()
     mock_credentials.token = mock_token
 
-    mock_default_auth.return_value = (mock_credentials, "project_id")
+    mock_default_auth.return_value = (mock_credentials, 'project_id')
 
     client_instance = OpenAIClient(
         location=mock_location,

py/plugins/vertex-ai/tests/model_garden/test_model_garden.py

@@ -16,3 +16,73 @@
 
 """Unittests for VertexAI Model Garden Models."""
 
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from genkit.ai import GenkitRegistry
+from genkit.plugins.compat_oai.typing import SupportedOutputFormat
+from genkit.plugins.vertex_ai.model_garden.model_garden import ModelGarden
+
+
+@pytest.fixture
+@patch('genkit.plugins.vertex_ai.model_garden.client.OpenAIClient')
+def model_garden_instance(client):
+    """Model Garden fixture."""
+    return ModelGarden(
+        model='test', location='us-central1', project_id='project', registry=MagicMock(spec=GenkitRegistry)
+    )
+
+
+@pytest.mark.parametrize(
+    'model_name, expected',
+    [
+        (
+            'meta/llama-3.1-405b-instruct-maas',
+            {
+                'name': 'ModelGarden - Meta - llama-3.1',
+                'supports': {
+                    'constrained': None,
+                    'content_type': None,
+                    'context': None,
+                    'multiturn': True,
+                    'media': False,
+                    'tools': True,
+                    'system_role': True,
+                    'output': [
+                        'json_mode',
+                        'text',
+                    ],
+                    'tool_choice': None,
+                },
+            },
+        ),
+        (
+            'meta/lazaro-model-pro-max',
+            {
+                'name': 'ModelGarden - meta/lazaro-model-pro-max',
+                'supports': {
+                    'constrained': None,
+                    'content_type': None,
+                    'context': None,
+                    'multiturn': True,
+                    'media': True,
+                    'tools': True,
+                    'system_role': True,
+                    'output': [
+                        'json_mode',
+                        'text',
+                    ],
+                    'tool_choice': None,
+                },
+            },
+        ),
+    ],
+)
+def test_get_model_info(model_name, expected, model_garden_instance):
+    """Unittest for get_model_info."""
+    model_garden_instance.name = model_name
+
+    result = model_garden_instance.get_model_info()
+
+    assert result == expected