MCP Server for Vertex AI Search
by ubie-oss
- tests
import unittest
from src.mcp_vertexai_search.config import (
Config,
DataStoreConfig,
GenerateContentConfig,
MCPServerConfig,
VertexAIModelConfig,
)
class TestConfig(unittest.TestCase):
def test_default_config(self):
"""Test that default Config values are set correctly."""
config = Config(
model=VertexAIModelConfig(
project_id="test-project",
model_name="test-model",
location="test-location",
),
)
self.assertIsInstance(config.server, MCPServerConfig)
self.assertIsInstance(config.model, VertexAIModelConfig)
self.assertEqual(config.data_stores, [])
def test_custom_config(self):
"""Test that Config can be initialized with custom values."""
custom_server = MCPServerConfig(name="test-server")
custom_model = VertexAIModelConfig(
project_id="test-project",
model_name="test-model",
location="test-location",
)
custom_data_store = DataStoreConfig(
project_id="test-project",
location="test-location",
datastore_id="test-datastore",
tool_name="test-tool",
)
config = Config(
server=custom_server,
model=custom_model,
data_stores=[custom_data_store]
)
self.assertEqual(config.server.name, "test-server")
self.assertEqual(config.model.model_name, "test-model")
self.assertEqual(config.model.location, "test-location")
self.assertEqual(len(config.data_stores), 1)
self.assertEqual(config.data_stores[0].datastore_id, "test-datastore")
def test_default_mcpserverconfig(self):
"""Test MCPServerConfig default values."""
server_config = MCPServerConfig()
self.assertEqual(server_config.name, "document-search")
def test_custom_mcpserverconfig(self):
"""Test MCPServerConfig with custom values."""
server_config = MCPServerConfig(name="custom-server")
self.assertEqual(server_config.name, "custom-server")
def test_default_vertexaimodelconfig(self):
"""Test VertexAIModelConfig default values."""
model_config = VertexAIModelConfig(
project_id="test-project",
location="test-location",
model_name="test-model",
)
self.assertIsInstance(model_config.generate_content_config, GenerateContentConfig)
self.assertEqual(model_config.project_id, "test-project")
self.assertEqual(model_config.location, "test-location")
self.assertEqual(model_config.model_name, "test-model")
self.assertEqual(model_config.generate_content_config.temperature, 0.7)
self.assertEqual(model_config.generate_content_config.top_p, 0.95)
def test_custom_vertexaimodelconfig(self):
"""Test VertexAIModelConfig with custom values."""
custom_gen_config = GenerateContentConfig(temperature=0.8, top_p=0.9)
model_config = VertexAIModelConfig(
model_name="custom-model",
location="custom-location",
project_id="custom-project",
generate_content_config=custom_gen_config,
)
self.assertEqual(model_config.model_name, "custom-model")
self.assertEqual(model_config.location, "custom-location")
self.assertEqual(model_config.project_id, "custom-project")
self.assertEqual(model_config.generate_content_config.temperature, 0.8)
self.assertEqual(model_config.generate_content_config.top_p, 0.9)
def test_default_generatecontentconfig(self):
"""Test GenerateContentConfig default values."""
gen_config = GenerateContentConfig()
self.assertEqual(gen_config.temperature, 0.7)
self.assertEqual(gen_config.top_p, 0.95)
def test_custom_generatecontentconfig(self):
"""Test GenerateContentConfig with custom values."""
gen_config = GenerateContentConfig(temperature=0.6, top_p=0.8)
self.assertEqual(gen_config.temperature, 0.6)
self.assertEqual(gen_config.top_p, 0.8)
def test_default_datastoreconfig(self):
"""Test DataStoreConfig default values."""
datastore_config = DataStoreConfig(
project_id="test-project",
location="test-location",
datastore_id="test-datastore",
tool_name="test-tool",
)
self.assertEqual(datastore_config.description, "")
self.assertEqual(datastore_config.tool_name, "test-tool")
def test_custom_datastoreconfig(self):
"""Test DataStoreConfig with custom values."""
datastore_config = DataStoreConfig(
project_id="custom-project",
location="custom-location",
datastore_id="custom-datastore",
description="custom-description",
tool_name="custom-tool",
)
self.assertEqual(datastore_config.project_id, "custom-project")
self.assertEqual(datastore_config.location, "custom-location")
self.assertEqual(datastore_config.datastore_id, "custom-datastore")
self.assertEqual(datastore_config.description, "custom-description")
self.assertEqual(datastore_config.tool_name, "custom-tool")
def test_computed_tool_name_datastoreconfig(self):
"""Test DataStoreConfig computed tool name when not provided."""
datastore_config = DataStoreConfig(
project_id="custom-project",
location="custom-location",
datastore_id="custom-datastore",
description="custom-description",
tool_name = "document-search"
)
expected_tool_name = "document-search"
self.assertEqual(datastore_config.tool_name, expected_tool_name)