fix unittestcase for fdp changes

Author: Pavan-Microsoft

src/api/agents/agent_factory.py

@@ -23,7 +23,7 @@ class AgentFactory:
     _lock = asyncio.Lock()
 
     @classmethod
-    async def get_instance(cls, config):
+    async def get_instance(cls):
         """
         Get or create the singleton AzureAIAgent instance.
         """

src/api/app.py

@@ -30,8 +30,7 @@ async def lifespan(fastapi_app: FastAPI):
     On startup, initializes the Azure AI agent using the configuration and attaches it to the app state.
     On shutdown, deletes the agent instance and performs any necessary cleanup.
     """
-    config = Config()
-    fastapi_app.state.agent = await AgentFactory.get_instance(config=config)
+    fastapi_app.state.agent = await AgentFactory.get_instance()
     yield
     await AgentFactory.delete_instance()
     fastapi_app.state.agent = None

src/tests/api/agents/test_agent_factory.py

@@ -12,24 +12,20 @@ def reset_agent_factory():
     AgentFactory._instance = None
 
 
-@pytest.fixture
-def config_mock():
-    return MagicMock(
-        azure_ai_project_conn_string="fake_conn_str",
-        azure_openai_deployment_model="gpt-4"
-    )
-
-
 @pytest.mark.asyncio
-@patch("agents.agent_factory.ChatWithDataPlugin", autospec=True)
+@patch("agents.agent_factory.AzureAIAgentSettings", autospec=True)
 @patch("agents.agent_factory.AzureAIAgent", autospec=True)
 @patch("agents.agent_factory.DefaultAzureCredential", autospec=True)
 async def test_get_instance_creates_new_agent(
     mock_credential,
     mock_azure_agent,
-    mock_plugin,
-    config_mock
+    mock_azure_ai_agent_settings
 ):
+    mock_settings_instance = MagicMock()
+    mock_settings_instance.endpoint = "https://fake-endpoint"
+    mock_settings_instance.model_deployment_name = "gpt-4o-mini"
+    mock_azure_ai_agent_settings.return_value = mock_settings_instance
+
     mock_client = AsyncMock()
     mock_agent_definition = MagicMock()
     mock_client.agents.create_agent.return_value = mock_agent_definition
@@ -38,21 +34,21 @@ async def test_get_instance_creates_new_agent(
     agent_instance = MagicMock()
     mock_azure_agent.return_value = agent_instance
 
-    result = await AgentFactory.get_instance(config_mock)
+    result = await AgentFactory.get_instance()
 
     assert result == agent_instance
     mock_azure_agent.create_client.assert_called_once_with(
         credential=mock_credential.return_value,
-        conn_str="fake_conn_str"
+        endpoint="https://fake-endpoint"
     )
     mock_client.agents.create_agent.assert_awaited_once()
     mock_azure_agent.assert_called_once()
 
 
 @pytest.mark.asyncio
-async def test_get_instance_returns_existing_instance(config_mock):
+async def test_get_instance_returns_existing_instance():
     AgentFactory._instance = MagicMock()
-    result = await AgentFactory.get_instance(config_mock)
+    result = await AgentFactory.get_instance()
     assert result == AgentFactory._instance
 
 

src/tests/api/plugins/test_chat_with_data_plugin.py

@@ -8,7 +8,6 @@ def mock_config():
     config_mock = MagicMock()
     config_mock.azure_openai_deployment_model = "gpt-4"
     config_mock.azure_openai_endpoint = "https://test-openai.azure.com/"
-    config_mock.azure_openai_api_key = "test-api-key"
     config_mock.azure_openai_api_version = "2024-02-15-preview"
     config_mock.azure_ai_search_endpoint = "https://search.test.azure.com/"
     config_mock.azure_ai_search_api_key = "search-api-key"
@@ -26,25 +25,29 @@ def chat_plugin(mock_config):
 
 
 class TestChatWithDataPlugin:
-    @pytest.mark.asyncio
+    @patch("plugins.chat_with_data_plugin.get_bearer_token_provider")
     @patch("plugins.chat_with_data_plugin.openai.AzureOpenAI")
-    async def test_greeting(self, mock_azure_openai, chat_plugin):
-        # Setup mock
+    @pytest.mark.asyncio
+    async def test_greeting(self, mock_azure_openai, mock_token_provider, chat_plugin):
+        # Setup mock token provider
+        mock_token_provider.return_value = lambda: "fake_token"
+        
+        # Setup mock client and completion response
         mock_client = MagicMock()
-        mock_azure_openai.return_value = mock_client
         mock_completion = MagicMock()
         mock_completion.choices = [MagicMock()]
         mock_completion.choices[0].message.content = "Hello, how can I help you?"
         mock_client.chat.completions.create.return_value = mock_completion
-        
+        mock_azure_openai.return_value = mock_client
+
         # Call the method
         result = await chat_plugin.greeting("Hello")
-        
+
         # Assertions
         assert result == "Hello, how can I help you?"
         mock_azure_openai.assert_called_once_with(
             azure_endpoint="https://test-openai.azure.com/",
-            api_key="test-api-key",
+            azure_ad_token_provider=mock_token_provider.return_value,
             api_version="2024-02-15-preview"
         )
         mock_client.chat.completions.create.assert_called_once()
@@ -103,10 +106,13 @@ async def test_greeting_exception(self, mock_azure_openai, chat_plugin):
         assert result == "Details could not be retrieved. Please try again later."
 
     @pytest.mark.asyncio
+    @patch("plugins.chat_with_data_plugin.get_bearer_token_provider")
     @patch("plugins.chat_with_data_plugin.execute_sql_query")
     @patch("plugins.chat_with_data_plugin.openai.AzureOpenAI")
-    async def test_get_SQL_Response(self, mock_azure_openai, mock_execute_sql, chat_plugin):
+    async def test_get_SQL_Response(self, mock_azure_openai, mock_execute_sql, mock_token_provider, chat_plugin):
+
         # Setup mocks
+        mock_token_provider.return_value = lambda: "fake_token"
         mock_client = MagicMock()
         mock_azure_openai.return_value = mock_client
         mock_completion = MagicMock()
@@ -123,7 +129,7 @@ async def test_get_SQL_Response(self, mock_azure_openai, mock_execute_sql, chat_
         assert result == "Query results data"
         mock_azure_openai.assert_called_once_with(
             azure_endpoint="https://test-openai.azure.com/",
-            api_key="test-api-key",
+            azure_ad_token_provider=mock_token_provider.return_value,
             api_version="2024-02-15-preview"
         )
         mock_client.chat.completions.create.assert_called_once()
@@ -175,9 +181,11 @@ async def test_get_SQL_Response_exception(self, mock_azure_openai, mock_execute_
         mock_execute_sql.assert_not_called()
 
     @pytest.mark.asyncio
+    @patch("plugins.chat_with_data_plugin.get_bearer_token_provider")
     @patch("plugins.chat_with_data_plugin.openai.AzureOpenAI")
-    async def test_get_answers_from_calltranscripts(self, mock_azure_openai, chat_plugin):
+    async def test_get_answers_from_calltranscripts(self, mock_azure_openai, mock_token_provider, chat_plugin):
         # Setup mock
+        mock_token_provider.return_value = lambda: "fake_token"
         mock_client = MagicMock()
         mock_azure_openai.return_value = mock_client
         mock_completion = MagicMock()
@@ -193,7 +201,7 @@ async def test_get_answers_from_calltranscripts(self, mock_azure_openai, chat_pl
         assert result.message.content == "Answer about transcripts"
         mock_azure_openai.assert_called_once_with(
             azure_endpoint="https://test-openai.azure.com/",
-            api_key="test-api-key",
+            azure_ad_token_provider=mock_token_provider.return_value,
             api_version="2024-02-15-preview"
         )
         mock_client.chat.completions.create.assert_called_once()

src/tests/api/services/test_chat_service.py

@@ -1,19 +1,17 @@
-import pytest
-from unittest.mock import AsyncMock, MagicMock, patch, call
 import json
-import uuid
 import time
-import asyncio
-import types
-from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
 from fastapi import HTTPException, status
 from semantic_kernel.exceptions.agent_exceptions import AgentException as RealAgentException
 
 
+
 # ---- Patch imports before importing the service under test ----
 @patch("common.config.config.Config")
 @patch("semantic_kernel.agents.AzureAIAgentThread")
-@patch("azure.ai.projects.models.TruncationObject")
+@patch("azure.ai.agents.models.TruncationObject")
 @patch("semantic_kernel.exceptions.agent_exceptions.AgentException")
 @patch("openai.AzureOpenAI")
 @patch("helpers.utils.format_stream_response")
@@ -23,7 +21,6 @@ def patched_imports(mock_format_stream, mock_openai, mock_agent_exception, mock_
     # Configure mock Config
     mock_config_instance = MagicMock()
     mock_config_instance.azure_openai_endpoint = "https://test.openai.azure.com"
-    mock_config_instance.azure_openai_api_key = "test_key"
     mock_config_instance.azure_openai_api_version = "2024-02-15-preview"
     mock_config_instance.azure_openai_deployment_model = "gpt-4o-mini"
     mock_config_instance.azure_ai_project_conn_string = "test_conn_string"
@@ -50,15 +47,14 @@ def patched_imports(mock_format_stream, mock_openai, mock_agent_exception, mock_
 # ---- Import service under test with patches ----
 with patch("common.config.config.Config") as mock_config, \
      patch("semantic_kernel.agents.AzureAIAgentThread") as mock_thread, \
-     patch("azure.ai.projects.models.TruncationObject") as mock_truncation, \
+     patch("azure.ai.agents.models.TruncationObject") as mock_truncation, \
      patch("semantic_kernel.exceptions.agent_exceptions.AgentException", new=RealAgentException) as mock_agent_exception, \
      patch("openai.AzureOpenAI") as mock_openai, \
      patch("helpers.utils.format_stream_response") as mock_format_stream:
 
     # Configure mock Config
     mock_config_instance = MagicMock()
     mock_config_instance.azure_openai_endpoint = "https://test.openai.azure.com"
-    mock_config_instance.azure_openai_api_key = "test_key"
     mock_config_instance.azure_openai_api_version = "2024-02-15-preview"
     mock_config_instance.azure_openai_deployment_model = "gpt-4o-mini"
     mock_config_instance.azure_ai_project_conn_string = "test_conn_string"
@@ -161,7 +157,6 @@ def test_init(self, mock_config_class, mock_request):
         # Configure mock Config
         mock_config_instance = MagicMock()
         mock_config_instance.azure_openai_endpoint = "https://test.openai.azure.com"
-        mock_config_instance.azure_openai_api_key = "test_key"
         mock_config_instance.azure_openai_api_version = "2024-02-15-preview"
         mock_config_instance.azure_openai_deployment_model = "gpt-4o-mini"
         mock_config_instance.azure_ai_project_conn_string = "test_conn_string"
@@ -173,10 +168,8 @@ def test_init(self, mock_config_class, mock_request):
         service = ChatService(mock_request)
 
         assert service.azure_openai_endpoint == "https://test.openai.azure.com"
-        assert service.azure_openai_api_key == "test_key"
         assert service.azure_openai_api_version == "2024-02-15-preview"
         assert service.azure_openai_deployment_name == "gpt-4o-mini"
-        assert service.azure_ai_project_conn_string == "test_conn_string"
         assert service.agent == mock_request.app.state.agent
         assert ChatService.thread_cache is not None