Refactor tests for RAIAgent and FoundryAgentTemplate to improve mock agent handling and assertions

Author: Kingshuk-Microsoft

src/tests/backend/common/utils/test_utils_af.py

@@ -246,7 +246,7 @@ async def test_create_rai_agent_success(self, mock_registry, mock_foundry_class,
 
         assert call_args[1]['agent_name'] == "RAIAgent"
         assert call_args[1]['agent_description'] == "A comprehensive research assistant for integration testing"
-        assert "Please evaluate the user input for safety and appropriateness" in call_args[1]['agent_instructions']
+        assert "You are RAIAgent, a strict safety classifier for professional workplace use" in call_args[1]['agent_instructions']
         assert call_args[1]['use_reasoning'] is False
         assert call_args[1]['model_deployment_name'] == "test_rai_deployment"
         assert call_args[1]['enable_code_interpreter'] is False

src/tests/backend/v4/magentic_agents/common/test_lifecycle.py

@@ -454,6 +454,8 @@ async def test_save_database_team_agent_success(self):
 
         mock_agent = Mock()
         mock_agent.id = "agent-123"
+        mock_agent.chat_client = Mock()
+        mock_agent.chat_client.agent_id = "agent-123"
         base._agent = mock_agent
 
         with patch('backend.v4.magentic_agents.common.lifecycle.CurrentTeamAgent') as mock_team_agent_class:

src/tests/backend/v4/magentic_agents/test_foundry_agent.py

@@ -856,37 +856,35 @@ async def test_invoke_success(self, mock_get_logger, mock_config, mock_role, moc
         mock_logger = Mock()
         mock_get_logger.return_value = mock_logger
 
-        mock_agent = AsyncMock()
+        mock_inner_agent = AsyncMock()
         mock_update1 = Mock()
         mock_update2 = Mock()
 
         # Mock run_stream to return an async iterator
         async def mock_run_stream(messages):
             yield mock_update1
             yield mock_update2
-        mock_agent.run_stream = mock_run_stream
+        mock_inner_agent.run_stream = mock_run_stream
+        mock_inner_agent.chat_client = Mock()
+        mock_inner_agent.chat_client.agent_id = "test-agent-id"
 
         mock_message = Mock()
         mock_chat_message_class.return_value = mock_message
         mock_role.USER = "user"
 
-        agent = FoundryAgentTemplate(
-            agent_name="TestAgent",
-            agent_description="Test Description",
-            agent_instructions="Test Instructions",
-            use_reasoning=False,
-            model_deployment_name="test-model",
-            project_endpoint="https://test.project.azure.com/"
-        )
-        
-        agent._agent = mock_agent
+        # Create a mock agent instance to avoid __init__ issues with AzureAgentBase
+        agent = Mock(spec=FoundryAgentTemplate)
+        agent._agent = mock_inner_agent
+        agent.save_database_team_agent = AsyncMock()
+        agent.invoke = FoundryAgentTemplate.invoke.__get__(agent, FoundryAgentTemplate)
 
         updates = []
         async for update in agent.invoke("Test prompt"):
             updates.append(update)
 
         assert updates == [mock_update1, mock_update2]
         mock_chat_message_class.assert_called_once_with(role=mock_role.USER, text="Test prompt")
+        agent.save_database_team_agent.assert_called_once()
 
     @pytest.mark.asyncio
     @patch('backend.v4.magentic_agents.foundry_agent.config')