Python: Improved AzureAI Package Test Coverage

python/packages/azure-ai/tests/test_agent_provider.py

@@ -19,6 +19,7 @@
     Agent,
     CodeInterpreterToolDefinition,
 )
+from azure.identity.aio import AzureCliCredential
 from pydantic import BaseModel
 
 from agent_framework_azure_ai import (
@@ -464,7 +465,77 @@ def test_as_agent_with_hosted_tools(
 
     assert isinstance(agent, ChatAgent)
     # Should have HostedCodeInterpreterTool in the default_options tools
-    assert any(isinstance(t, HostedCodeInterpreterTool) for t in (agent.default_options.get("tools") or []))
+    assert any(isinstance(t, HostedCodeInterpreterTool) for t in (agent.default_options.get("tools") or []))  # type: ignore
+
+
+def test_as_agent_with_dict_function_tools_validates(
+    azure_ai_unit_test_env: dict[str, str],
+    mock_agents_client: MagicMock,
+) -> None:
+    """Test as_agent validates dict-format function tools require implementations."""
+    # Dict-based function tool (as returned by some Azure AI SDK operations)
+    dict_function_tool = {  # type: ignore
+        "type": "function",
+        "function": {
+            "name": "dict_based_function",
+            "description": "A function defined as dict",
+            "parameters": {"type": "object", "properties": {}},
+        },
+    }
+
+    mock_agent = MagicMock(spec=Agent)
+    mock_agent.id = "agent-id"
+    mock_agent.name = "Agent"
+    mock_agent.description = None
+    mock_agent.instructions = None
+    mock_agent.model = "gpt-4"
+    mock_agent.temperature = None
+    mock_agent.top_p = None
+    mock_agent.tools = [dict_function_tool]
+
+    provider = AzureAIAgentsProvider(agents_client=mock_agents_client)
+
+    with pytest.raises(ServiceInitializationError) as exc_info:
+        provider.as_agent(mock_agent)
+
+    assert "dict_based_function" in str(exc_info.value)
+
+
+def test_as_agent_with_dict_function_tools_provided(
+    azure_ai_unit_test_env: dict[str, str],
+    mock_agents_client: MagicMock,
+) -> None:
+    """Test as_agent succeeds when dict-format function tools have implementations provided."""
+    dict_function_tool = {  # type: ignore
+        "type": "function",
+        "function": {
+            "name": "dict_based_function",
+            "description": "A function defined as dict",
+            "parameters": {"type": "object", "properties": {}},
+        },
+    }
+
+    mock_agent = MagicMock(spec=Agent)
+    mock_agent.id = "agent-id"
+    mock_agent.name = "Agent"
+    mock_agent.description = None
+    mock_agent.instructions = None
+    mock_agent.model = "gpt-4"
+    mock_agent.temperature = None
+    mock_agent.top_p = None
+    mock_agent.tools = [dict_function_tool]
+
+    @tool
+    def dict_based_function() -> str:
+        """A function implementation."""
+        return "result"
+
+    provider = AzureAIAgentsProvider(agents_client=mock_agents_client)
+
+    agent = provider.as_agent(mock_agent, tools=dict_based_function)
+
+    assert isinstance(agent, ChatAgent)
+    assert agent.id == "agent-id"
 
 
 # endregion
@@ -729,8 +800,6 @@ def test_from_azure_ai_agent_tools_unknown_dict() -> None:
 @skip_if_azure_ai_integration_tests_disabled
 async def test_integration_create_agent() -> None:
     """Integration test: Create an agent using the provider."""
-    from azure.identity.aio import AzureCliCredential
-
     async with (
         AzureCliCredential() as credential,
         AzureAIAgentsProvider(credential=credential) as provider,
@@ -753,8 +822,6 @@ async def test_integration_create_agent() -> None:
 @skip_if_azure_ai_integration_tests_disabled
 async def test_integration_get_agent() -> None:
     """Integration test: Get an existing agent using the provider."""
-    from azure.identity.aio import AzureCliCredential
-
     async with (
         AzureCliCredential() as credential,
         AzureAIAgentsProvider(credential=credential) as provider,
@@ -779,8 +846,6 @@ async def test_integration_get_agent() -> None:
 @skip_if_azure_ai_integration_tests_disabled
 async def test_integration_create_and_run() -> None:
     """Integration test: Create an agent and run a conversation."""
-    from azure.identity.aio import AzureCliCredential
-
     async with (
         AzureCliCredential() as credential,
         AzureAIAgentsProvider(credential=credential) as provider,

python/packages/azure-ai/tests/test_azure_ai_agent_client.py

@@ -26,16 +26,19 @@
     tool,
 )
 from agent_framework._serialization import SerializationMixin
-from agent_framework.exceptions import ServiceInitializationError
+from agent_framework.exceptions import ServiceInitializationError, ServiceInvalidRequestError
 from azure.ai.agents.models import (
     AgentsNamedToolChoice,
     AgentsNamedToolChoiceType,
+    AgentsToolChoiceOptionMode,
+    CodeInterpreterToolDefinition,
     FileInfo,
     MessageDeltaChunk,
     MessageDeltaTextContent,
     MessageDeltaTextFileCitationAnnotation,
     MessageDeltaTextFilePathAnnotation,
     MessageDeltaTextUrlCitationAnnotation,
+    MessageInputTextBlock,
     RequiredFunctionToolCall,
     RequiredMcpToolCall,
     RunStatus,
@@ -593,8 +596,6 @@ async def test_azure_ai_chat_client_prepare_options_with_none_tool_choice(
 
     run_options, _ = await chat_client._prepare_options([], chat_options)  # type: ignore
 
-    from azure.ai.agents.models import AgentsToolChoiceOptionMode
-
     assert run_options["tool_choice"] == AgentsToolChoiceOptionMode.NONE
 
 
@@ -608,8 +609,6 @@ async def test_azure_ai_chat_client_prepare_options_with_auto_tool_choice(
 
     run_options, _ = await chat_client._prepare_options([], chat_options)  # type: ignore
 
-    from azure.ai.agents.models import AgentsToolChoiceOptionMode
-
     assert run_options["tool_choice"] == AgentsToolChoiceOptionMode.AUTO
 
 
@@ -1942,3 +1941,205 @@ def test_azure_ai_chat_client_init_with_auto_created_agents_client(
         assert client.agent_id == "test-agent"
         assert client.credential is mock_azure_credential
         assert client._should_close_client is True  # Should close since we created it  # type: ignore[attr-defined]
+
+
+async def test_azure_ai_chat_client_prepare_options_with_mapping_response_format(
+    mock_agents_client: MagicMock,
+) -> None:
+    """Test _prepare_options with Mapping-based response_format (runtime JSON schema)."""
+    chat_client = create_test_azure_ai_chat_client(mock_agents_client)
+
+    # Runtime JSON schema dict
+    response_format_dict = {
+        "type": "json_schema",
+        "json_schema": {
+            "name": "TestSchema",
+            "schema": {"type": "object", "properties": {"name": {"type": "string"}}},
+        },
+    }
+
+    chat_options: ChatOptions = {"response_format": response_format_dict}  # type: ignore[typeddict-item]
+
+    run_options, _ = await chat_client._prepare_options([], chat_options)  # type: ignore
+
+    assert "response_format" in run_options
+    # Should pass through as-is for Mapping types
+    assert run_options["response_format"] == response_format_dict
+
+
+async def test_azure_ai_chat_client_prepare_options_with_invalid_response_format(
+    mock_agents_client: MagicMock,
+) -> None:
+    """Test _prepare_options with invalid response_format raises error."""
+    chat_client = create_test_azure_ai_chat_client(mock_agents_client)
+
+    # Invalid response_format (not BaseModel or Mapping)
+    chat_options: ChatOptions = {"response_format": "invalid_format"}  # type: ignore[typeddict-item]
+
+    with pytest.raises(ServiceInvalidRequestError, match="response_format must be a Pydantic BaseModel"):
+        await chat_client._prepare_options([], chat_options)  # type: ignore
+
+
+async def test_azure_ai_chat_client_prepare_tool_definitions_with_agent_tool_resources(
+    mock_agents_client: MagicMock,
+) -> None:
+    """Test _prepare_tool_definitions_and_resources copies tool_resources from agent definition."""
+    chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
+
+    # Create mock agent definition with tool_resources
+    mock_agent_definition = MagicMock()
+    mock_agent_definition.tools = []
+    mock_agent_definition.tool_resources = {"code_interpreter": {"file_ids": ["file-123"]}}
+
+    run_options: dict[str, Any] = {}
+    options: dict[str, Any] = {}
+
+    await chat_client._prepare_tool_definitions_and_resources(options, mock_agent_definition, run_options)  # type: ignore
+
+    # Verify tool_resources was copied to run_options
+    assert "tool_resources" in run_options
+    assert run_options["tool_resources"] == {"code_interpreter": {"file_ids": ["file-123"]}}
+
+
+def test_azure_ai_chat_client_prepare_mcp_resources_with_dict_approval_mode(
+    mock_agents_client: MagicMock,
+) -> None:
+    """Test _prepare_mcp_resources with dict-based approval mode (always_require_approval)."""
+    chat_client = create_test_azure_ai_chat_client(mock_agents_client)
+
+    # MCP tool with dict-based approval mode
+    mcp_tool = HostedMCPTool(
+        name="Test MCP",
+        url="https://example.com/mcp",
+        approval_mode={"always_require_approval": {"tool1", "tool2"}},
+    )
+
+    result = chat_client._prepare_mcp_resources([mcp_tool])  # type: ignore
+
+    assert len(result) == 1
+    assert result[0]["server_label"] == "Test_MCP"
+    assert "require_approval" in result[0]
+    assert result[0]["require_approval"] == {"always": {"tool1", "tool2"}}
+
+
+def test_azure_ai_chat_client_prepare_mcp_resources_with_never_require_dict(
+    mock_agents_client: MagicMock,
+) -> None:
+    """Test _prepare_mcp_resources with dict-based approval mode (never_require_approval)."""
+    chat_client = create_test_azure_ai_chat_client(mock_agents_client)
+
+    # MCP tool with never_require_approval dict
+    mcp_tool = HostedMCPTool(
+        name="Test MCP",
+        url="https://example.com/mcp",
+        approval_mode={"never_require_approval": {"safe_tool"}},
+    )
+
+    result = chat_client._prepare_mcp_resources([mcp_tool])  # type: ignore
+
+    assert len(result) == 1
+    assert result[0]["require_approval"] == {"never": {"safe_tool"}}
+
+
+def test_azure_ai_chat_client_prepare_messages_with_function_result(
+    mock_agents_client: MagicMock,
+) -> None:
+    """Test _prepare_messages extracts function_result content."""
+    chat_client = create_test_azure_ai_chat_client(mock_agents_client)
+
+    function_result = Content.from_function_result(call_id='["run_123", "call_456"]', result="test result")
+    messages = [ChatMessage(role=Role.USER, contents=[function_result])]
+
+    additional_messages, instructions, required_action_results = chat_client._prepare_messages(messages)  # type: ignore
+
+    # function_result should be extracted, not added to additional_messages
+    assert additional_messages is None
+    assert required_action_results is not None
+    assert len(required_action_results) == 1
+    assert required_action_results[0].type == "function_result"
+
+
+def test_azure_ai_chat_client_prepare_messages_with_raw_content_block(
+    mock_agents_client: MagicMock,
+) -> None:
+    """Test _prepare_messages handles raw MessageInputContentBlock in content."""
+    chat_client = create_test_azure_ai_chat_client(mock_agents_client)
+
+    # Create content with raw_representation that is a MessageInputContentBlock
+    raw_block = MessageInputTextBlock(text="Raw block text")
+    custom_content = Content(type="custom", raw_representation=raw_block)
+    messages = [ChatMessage(role=Role.USER, contents=[custom_content])]
+
+    additional_messages, instructions, required_action_results = chat_client._prepare_messages(messages)  # type: ignore
+
+    assert additional_messages is not None
+    assert len(additional_messages) == 1
+    assert len(additional_messages[0].content) == 1
+    assert additional_messages[0].content[0] == raw_block
+
+
+async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_mcp_tool(
+    mock_agents_client: MagicMock,
+) -> None:
+    """Test _prepare_tools_for_azure_ai with HostedMCPTool."""
+    chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
+
+    mcp_tool = HostedMCPTool(
+        name="Test MCP Server",
+        url="https://example.com/mcp",
+        allowed_tools=["tool1", "tool2"],
+    )
+
+    tool_definitions = await chat_client._prepare_tools_for_azure_ai([mcp_tool])  # type: ignore
+
+    assert len(tool_definitions) >= 1
+    # The McpTool.definitions property returns the tool definitions
+    # Verify the MCP tool was converted correctly by checking the definition type
+    mcp_def = tool_definitions[0]
+    assert mcp_def.get("type") == "mcp"
+
+
+async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_tool_definition(
+    mock_agents_client: MagicMock,
+) -> None:
+    """Test _prepare_tools_for_azure_ai with ToolDefinition passthrough."""
+    chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
+
+    # Pass a ToolDefinition directly - should be passed through as-is
+    tool_def = CodeInterpreterToolDefinition()
+
+    tool_definitions = await chat_client._prepare_tools_for_azure_ai([tool_def])  # type: ignore
+
+    assert len(tool_definitions) == 1
+    assert tool_definitions[0] is tool_def
+
+
+async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_dict_passthrough(
+    mock_agents_client: MagicMock,
+) -> None:
+    """Test _prepare_tools_for_azure_ai with dict passthrough."""
+    chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
+
+    # Pass a dict tool definition - should be passed through as-is
+    dict_tool = {"type": "function", "function": {"name": "test_func", "parameters": {}}}
+
+    tool_definitions = await chat_client._prepare_tools_for_azure_ai([dict_tool])  # type: ignore
+
+    assert len(tool_definitions) == 1
+    assert tool_definitions[0] is dict_tool
+
+
+async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_unsupported_type(
+    mock_agents_client: MagicMock,
+) -> None:
+    """Test _prepare_tools_for_azure_ai raises error for unsupported tool type."""
+    chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
+
+    # Pass an unsupported tool type
+    class UnsupportedTool:
+        pass
+
+    unsupported_tool = UnsupportedTool()
+
+    with pytest.raises(ServiceInitializationError, match="Unsupported tool type"):
+        await chat_client._prepare_tools_for_azure_ai([unsupported_tool])  # type: ignore