unittestcase changes for sqlagent

Author: Pavan-Microsoft

src/tests/api/agents/test_base_agent_factory.py

@@ -0,0 +1,90 @@
+import pytest
+import asyncio
+from unittest.mock import AsyncMock, patch, MagicMock
+
+from common.config.config import Config
+from agents.agent_factory_base import BaseAgentFactory
+
+
+class MockAgentFactory(BaseAgentFactory):
+    """Concrete test class extending BaseAgentFactory for unit testing."""
+    _created = False
+    _deleted = False
+
+    @classmethod
+    async def create_agent(cls, config: Config):
+        cls._created = True
+        return {"agent": "mock-agent"}
+
+    @classmethod
+    async def _delete_agent_instance(cls, agent: object):
+        cls._deleted = True
+
+
+@pytest.fixture(autouse=True)
+def reset_factory_state():
+    MockAgentFactory._agent = None
+    MockAgentFactory._created = False
+    MockAgentFactory._deleted = False
+    yield
+    MockAgentFactory._agent = None
+
+
+@pytest.mark.asyncio
+async def test_get_agent_creates_singleton():
+    # Agent should be None initially
+    assert MockAgentFactory._agent is None
+
+    result1 = await MockAgentFactory.get_agent()
+    result2 = await MockAgentFactory.get_agent()
+
+    # Should be the same object
+    assert result1 is result2
+    assert MockAgentFactory._created is True
+    assert MockAgentFactory._agent == {"agent": "mock-agent"}
+
+
+@pytest.mark.asyncio
+async def test_delete_agent_removes_singleton():
+    # Set initial agent
+    await MockAgentFactory.get_agent()
+    assert MockAgentFactory._agent is not None
+
+    await MockAgentFactory.delete_agent()
+
+    assert MockAgentFactory._agent is None
+    assert MockAgentFactory._deleted is True
+
+
+@pytest.mark.asyncio
+async def test_delete_agent_does_nothing_if_none():
+    # Agent is None
+    await MockAgentFactory.delete_agent()
+
+    assert MockAgentFactory._agent is None
+    assert MockAgentFactory._deleted is False
+
+
+@pytest.mark.asyncio
+async def test_thread_safety_of_get_agent(monkeypatch):
+    # Patch create_agent to delay and track calls
+    call_count = 0
+
+    async def slow_create_agent(config):
+        nonlocal call_count
+        call_count += 1
+        await asyncio.sleep(0.1)
+        return {"agent": "thread-safe"}
+
+    monkeypatch.setattr(MockAgentFactory, "create_agent", slow_create_agent)
+
+    # Run get_agent concurrently
+    results = await asyncio.gather(
+        MockAgentFactory.get_agent(),
+        MockAgentFactory.get_agent(),
+        MockAgentFactory.get_agent()
+    )
+
+    # All should return the same instance
+    assert all(result == {"agent": "thread-safe"} for result in results)
+    assert call_count == 1  # Only one creation

src/tests/api/agents/test_search_agent_factory.py

@@ -1,7 +1,5 @@
 import pytest
-import asyncio
-from unittest.mock import patch, MagicMock, AsyncMock
-
+from unittest.mock import patch, MagicMock
 from agents.search_agent_factory import SearchAgentFactory
 
 
@@ -13,48 +11,49 @@ def reset_search_agent_factory():
 
 
 @pytest.mark.asyncio
-@patch("agents.search_agent_factory.Config", autospec=True)
 @patch("agents.search_agent_factory.DefaultAzureCredential", autospec=True)
 @patch("agents.search_agent_factory.AIProjectClient", autospec=True)
 @patch("agents.search_agent_factory.AzureAISearchTool", autospec=True)
-async def test_get_agent_creates_new_instance(
-    mock_search_tool,
-    mock_project_client_class,
-    mock_credential,
-    mock_config_class
+async def test_create_agent_creates_new_instance(
+    mock_search_tool_cls,
+    mock_project_client_cls,
+    mock_credential_cls
 ):
     # Mock config
     mock_config = MagicMock()
     mock_config.ai_project_endpoint = "https://fake-endpoint"
     mock_config.azure_ai_search_connection_name = "fake-connection"
     mock_config.azure_ai_search_index = "fake-index"
     mock_config.azure_openai_deployment_model = "fake-model"
-    mock_config_class.return_value = mock_config
+    mock_config.solution_name = "test-solution"
+    mock_config.ai_project_api_version = "2025-05-01"
 
     # Mock project client
     mock_project_client = MagicMock()
-    mock_project_client_class.return_value = mock_project_client
+    mock_project_client_cls.return_value = mock_project_client
 
+    # Mock index response
     mock_index = MagicMock()
     mock_index.name = "index-name"
     mock_index.version = "1"
     mock_project_client.indexes.create_or_update.return_value = mock_index
 
     # Mock search tool
-    mock_tool = MagicMock()
-    mock_tool.definitions = ["tool-def"]
-    mock_tool.resources = ["tool-res"]
-    mock_search_tool.return_value = mock_tool
+    mock_search_tool_instance = MagicMock()
+    mock_search_tool_instance.definitions = ["tool-def"]
+    mock_search_tool_instance.resources = ["tool-res"]
+    mock_search_tool_cls.return_value = mock_search_tool_instance
 
     # Mock agent
     mock_agent = MagicMock()
     mock_project_client.agents.create_agent.return_value = mock_agent
 
-    # Run the factory
-    result = await SearchAgentFactory.get_agent()
+    # Run the factory directly
+    result = await SearchAgentFactory.create_agent(mock_config)
 
     assert result["agent"] == mock_agent
     assert result["client"] == mock_project_client
+
     mock_project_client.indexes.create_or_update.assert_called_once_with(
         name="project-index-fake-connection-fake-index",
         version="1",
@@ -74,23 +73,20 @@ async def test_get_agent_creates_new_instance(
 
 @pytest.mark.asyncio
 async def test_get_agent_returns_existing_instance():
+    # Setup: Already initialized
     SearchAgentFactory._agent = {"agent": MagicMock(), "client": MagicMock()}
     result = await SearchAgentFactory.get_agent()
     assert result == SearchAgentFactory._agent
 
 
 @pytest.mark.asyncio
 async def test_delete_agent_removes_agent():
+    # Setup
     mock_agent = MagicMock()
     mock_agent.id = "mock-agent-id"
-
     mock_client = MagicMock()
-    mock_client.agents.delete_agent = MagicMock()
 
-    SearchAgentFactory._agent = {
-        "agent": mock_agent,
-        "client": mock_client
-    }
+    SearchAgentFactory._agent = {"agent": mock_agent, "client": mock_client}
 
     await SearchAgentFactory.delete_agent()
 
@@ -99,6 +95,8 @@ async def test_delete_agent_removes_agent():
 
 
 @pytest.mark.asyncio
-def test_delete_agent_does_nothing_if_none():
+async def test_delete_agent_does_nothing_if_none():
     SearchAgentFactory._agent = None
-    SearchAgentFactory.delete_agent()
+    await SearchAgentFactory.delete_agent()
+    # No error should be raised, and nothing is called
+

src/tests/api/agents/test_sql_agent_factory.py

@@ -0,0 +1,83 @@
+import pytest
+from unittest.mock import patch, MagicMock, AsyncMock
+
+from agents.sql_agent_factory import SQLAgentFactory
+
+
+@pytest.fixture(autouse=True)
+def reset_sql_agent_factory():
+    SQLAgentFactory._agent = None
+    yield
+    SQLAgentFactory._agent = None
+
+
+@pytest.mark.asyncio
+@patch("agents.sql_agent_factory.DefaultAzureCredential", autospec=True)
+@patch("agents.sql_agent_factory.AIProjectClient", autospec=True)
+async def test_create_agent_creates_new_instance(
+    mock_ai_client_cls,
+    mock_credential_cls
+):
+    # Mock config
+    mock_config = MagicMock()
+    mock_config.ai_project_endpoint = "https://test-endpoint"
+    mock_config.ai_project_api_version = "2025-05-01"
+    mock_config.azure_openai_deployment_model = "test-model"
+    mock_config.solution_name = "test-solution"
+
+    # Mock project client
+    mock_project_client = MagicMock()
+    mock_ai_client_cls.return_value = mock_project_client
+
+    # Mock agent
+    mock_agent = MagicMock()
+    mock_project_client.agents.create_agent.return_value = mock_agent
+
+    result = await SQLAgentFactory.create_agent(mock_config)
+
+    assert result["agent"] == mock_agent
+    assert result["client"] == mock_project_client
+
+    mock_ai_client_cls.assert_called_once_with(
+        endpoint="https://test-endpoint",
+        credential=mock_credential_cls.return_value,
+        api_version="2025-05-01"
+    )
+    mock_project_client.agents.create_agent.assert_called_once()
+    args, kwargs = mock_project_client.agents.create_agent.call_args
+    assert kwargs["model"] == "test-model"
+    assert kwargs["name"] == "KM-ChatWithSQLDatabaseAgent-test-solution"
+    assert "Generate a valid T-SQL query" in kwargs["instructions"]
+
+
+@pytest.mark.asyncio
+async def test_get_agent_returns_existing_instance():
+    SQLAgentFactory._agent = {"agent": MagicMock(), "client": MagicMock()}
+    result = await SQLAgentFactory.get_agent()
+    assert result == SQLAgentFactory._agent
+
+
+@pytest.mark.asyncio
+async def test_delete_agent_removes_agent():
+    mock_agent = MagicMock()
+    mock_agent.id = "agent-id"
+
+    mock_client = MagicMock()
+    mock_client.agents.delete_agent = MagicMock()
+
+    SQLAgentFactory._agent = {
+        "agent": mock_agent,
+        "client": mock_client
+    }
+
+    await SQLAgentFactory.delete_agent()
+
+    mock_client.agents.delete_agent.assert_called_once_with("agent-id")
+    assert SQLAgentFactory._agent is None
+
+
+@pytest.mark.asyncio
+async def test_delete_agent_does_nothing_if_none():
+    SQLAgentFactory._agent = None
+    await SQLAgentFactory.delete_agent()
+    # Nothing should raise, nothing should be called