chore: import upstream snapshot with attribution

python/packages/a2a/tests/test_a2a_executor.py

@@ -0,0 +1,907 @@
+# Copyright (c) Microsoft. All rights reserved.
+from asyncio import CancelledError
+from typing import Any, cast
+from unittest.mock import AsyncMock, MagicMock, patch
+from uuid import uuid4
+
+from a2a.types import Part, Task, TaskState
+from agent_framework import (
+    AgentResponseUpdate,
+    Content,
+    Message,
+    SupportsAgentRun,
+)
+from agent_framework._types import AgentResponse
+from agent_framework.a2a import A2AExecutor
+from pytest import fixture, raises
+
+
+@fixture
+def mock_agent() -> MagicMock:
+    """Fixture that provides a mock SupportsAgentRun."""
+    agent = MagicMock(spec=SupportsAgentRun)
+    agent.run = AsyncMock()
+    return agent
+
+
+@fixture
+def mock_request_context() -> MagicMock:
+    """Fixture that provides a mock RequestContext."""
+    request_context = MagicMock()
+    request_context.context_id = str(uuid4())
+    request_context.get_user_input = MagicMock(return_value="Test query")
+    request_context.current_task = None
+    request_context.message = None
+    return request_context
+
+
+@fixture
+def mock_event_queue() -> MagicMock:
+    """Fixture that provides a mock EventQueue."""
+    queue = AsyncMock()
+    queue.enqueue_event = AsyncMock()
+    return queue
+
+
+@fixture
+def mock_task() -> Task:
+    """Fixture that provides a mock Task."""
+    task = MagicMock(spec=Task)
+    task.id = str(uuid4())
+    task.context_id = str(uuid4())
+    task.state = TaskState.TASK_STATE_COMPLETED
+    return task
+
+
+@fixture
+def mock_task_updater() -> MagicMock:
+    """Fixture that provides a mock TaskUpdater."""
+    updater = MagicMock()
+    updater.submit = AsyncMock()
+    updater.start_work = AsyncMock()
+    updater.complete = AsyncMock()
+    updater.update_status = AsyncMock()
+    updater.new_agent_message = MagicMock()
+    return updater
+
+
+@fixture
+def executor(mock_agent: MagicMock) -> A2AExecutor:
+    """Fixture that provides an A2AExecutor."""
+    return A2AExecutor(agent=mock_agent)
+
+
+class TestA2AExecutorInitialization:
+    """Tests for A2AExecutor initialization."""
+
+    def test_initialization_with_agent_only(self, mock_agent: MagicMock) -> None:
+        """Arrange: Create mock agent
+        Act: Initialize A2AExecutor with only agent
+        Assert: Executor is created with default values
+        """
+        # Act
+        executor = A2AExecutor(agent=mock_agent)
+
+        # Assert
+        assert executor._agent is mock_agent
+        assert executor._stream is False
+        assert executor._run_kwargs == {}
+
+    def test_initialization_with_stream_and_kwargs(self, mock_agent: MagicMock) -> None:
+        """Arrange: Create mock agent
+        Act: Initialize A2AExecutor with stream and run_kwargs
+        Assert: Executor is created with specified values
+        """
+        # Arrange
+        run_kwargs = {"temperature": 0.5}
+
+        # Act
+        executor = A2AExecutor(agent=mock_agent, stream=True, run_kwargs=run_kwargs)
+
+        # Assert
+        assert executor._agent is mock_agent
+        assert executor._stream is True
+        assert executor._run_kwargs == run_kwargs
+
+    def test_initialization_with_invalid_run_kwargs(self, mock_agent: MagicMock) -> None:
+        """Arrange: Create mock agent
+        Act: Initialize A2AExecutor with reserved keys in run_kwargs
+        Assert: ValueError is raised
+        """
+        # Act & Assert
+        with raises(ValueError, match="run_kwargs cannot contain 'session'"):
+            A2AExecutor(agent=mock_agent, run_kwargs={"session": "something"})
+
+        with raises(ValueError, match="run_kwargs cannot contain 'stream'"):
+            A2AExecutor(agent=mock_agent, run_kwargs={"stream": True})
+
+
+class TestA2AExecutorCancel:
+    """Tests for the cancel method."""
+
+    async def test_cancel_method_completes(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+    ) -> None:
+        """Arrange: Create executor with dependencies
+        Act: Call cancel method
+        Assert: Method completes without raising error
+        """
+        # Arrange
+        mock_request_context.task_id = "task-123"
+
+        # Act & Assert (should not raise)
+        await executor.cancel(mock_request_context, mock_event_queue)  # type: ignore
+
+    async def test_cancel_handles_different_contexts(
+        self,
+        executor: A2AExecutor,
+        mock_event_queue: MagicMock,
+    ) -> None:
+        """Arrange: Create executor with multiple request contexts
+        Act: Call cancel with different contexts
+        Assert: Each cancel completes successfully
+        """
+        # Arrange
+        context1 = MagicMock()
+        context1.context_id = "ctx-1"
+        context1.task_id = "task-1"
+        context2 = MagicMock()
+        context2.context_id = "ctx-2"
+        context2.task_id = "task-2"
+
+        # Act & Assert
+        await executor.cancel(context1, mock_event_queue)  # type: ignore
+        await executor.cancel(context2, mock_event_queue)  # type: ignore
+
+    async def test_cancel_raises_error_when_context_id_missing(
+        self,
+        executor: A2AExecutor,
+        mock_event_queue: MagicMock,
+    ) -> None:
+        """Arrange: Create context without context_id
+        Act: Call cancel method
+        Assert: ValueError is raised
+        """
+        # Arrange
+        mock_context = MagicMock()
+        mock_context.context_id = None
+
+        # Act & Assert
+        with raises(ValueError) as excinfo:
+            await executor.cancel(mock_context, mock_event_queue)  # type: ignore
+
+        # Assert
+        assert "Context ID" in str(excinfo.value)
+
+
+class TestA2AExecutorExecute:
+    """Tests for the execute method."""
+
+    async def test_execute_with_existing_task_succeeds(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+        mock_task: Task,
+    ) -> None:
+        """Arrange: Create executor with mocked dependencies and existing task
+        Act: Call execute method
+        Assert: Execution completes successfully
+        """
+        # Arrange
+        mock_request_context.get_user_input = MagicMock(return_value="Hello")
+        mock_request_context.current_task = mock_task
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = MagicMock()
+
+        response_message = Message(role="assistant", contents=[Content.from_text(text="Hello back")])
+        response = MagicMock(spec=AgentResponse)
+        response.messages = [response_message]
+        cast(Any, executor._agent).run = AsyncMock(return_value=response)
+        cast(Any, executor._agent).create_session = MagicMock()
+
+        with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+            mock_updater = MagicMock()
+            mock_updater.submit = AsyncMock()
+            mock_updater.start_work = AsyncMock()
+            mock_updater.complete = AsyncMock()
+            mock_updater.update_status = AsyncMock()
+            mock_updater.new_agent_message = MagicMock(return_value="message_obj")
+            mock_updater_class.return_value = mock_updater
+
+            # Act
+            await executor.execute(mock_request_context, mock_event_queue)
+
+            # Assert
+            mock_updater.submit.assert_called_once()
+            mock_updater.start_work.assert_called_once()
+            mock_updater.complete.assert_called_once()
+            cast(Any, executor._agent.create_session).assert_called_once()
+            cast(Any, executor._agent.run).assert_called_once()
+
+    async def test_execute_creates_task_when_not_exists(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+    ) -> None:
+        """Arrange: Create executor with request context without task
+        Act: Call execute method
+        Assert: New task is created and enqueued
+        """
+        # Arrange
+        mock_message = MagicMock()
+        mock_request_context.get_user_input = MagicMock(return_value="Hello")
+        mock_request_context.current_task = None
+        mock_request_context.message = mock_message
+        mock_request_context.context_id = "ctx-123"
+
+        response_message = Message(role="assistant", contents=[Content.from_text(text="Response")])
+        response = MagicMock(spec=AgentResponse)
+        response.messages = [response_message]
+        cast(Any, executor._agent).run = AsyncMock(return_value=response)
+        cast(Any, executor._agent).create_session = MagicMock()
+
+        with patch("agent_framework_a2a._a2a_executor.new_task_from_user_message") as mock_new_task:
+            mock_task = MagicMock(spec=Task)
+            mock_task.id = "task-new"
+            mock_task.context_id = "ctx-123"
+            mock_new_task.return_value = mock_task
+
+            with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+                mock_updater = MagicMock()
+                mock_updater.submit = AsyncMock()
+                mock_updater.start_work = AsyncMock()
+                mock_updater.complete = AsyncMock()
+                mock_updater.update_status = AsyncMock()
+                mock_updater.new_agent_message = MagicMock(return_value="message_obj")
+                mock_updater_class.return_value = mock_updater
+
+                # Act
+                await executor.execute(mock_request_context, mock_event_queue)
+
+                # Assert
+                mock_new_task.assert_called_once()
+                mock_event_queue.enqueue_event.assert_called_once()
+
+    async def test_execute_raises_error_when_context_id_missing(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+    ) -> None:
+        """Arrange: Create context without context_id
+        Act: Call execute method
+        Assert: ValueError is raised
+        """
+        # Arrange
+        mock_request_context.context_id = None
+        mock_request_context.message = MagicMock()
+
+        # Act & Assert
+        with raises(ValueError) as excinfo:
+            await executor.execute(mock_request_context, mock_event_queue)
+
+        # Assert
+        assert "Context ID" in str(excinfo.value)
+
+    async def test_execute_raises_error_when_message_missing(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+    ) -> None:
+        """Arrange: Create context without message
+        Act: Call execute method
+        Assert: ValueError is raised
+        """
+        # Arrange
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = None
+
+        # Act & Assert
+        with raises(ValueError) as excinfo:
+            await executor.execute(mock_request_context, mock_event_queue)
+
+        # Assert
+        assert "Message" in str(excinfo.value)
+
+    async def test_execute_handles_cancelled_error(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+        mock_task: Task,
+    ) -> None:
+        """Arrange: Create executor that raises CancelledError
+        Act: Call execute method
+        Assert: Error is caught and task is marked as canceled
+        """
+        # Arrange
+        mock_request_context.get_user_input = MagicMock(return_value="Hello")
+        mock_request_context.current_task = mock_task
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = MagicMock()
+
+        cast(Any, executor._agent).run = AsyncMock(side_effect=CancelledError())
+        cast(Any, executor._agent).create_session = MagicMock()
+
+        with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+            mock_updater = MagicMock()
+            mock_updater.submit = AsyncMock()
+            mock_updater.start_work = AsyncMock()
+            mock_updater.update_status = AsyncMock()
+            mock_updater_class.return_value = mock_updater
+
+            # Act
+            await executor.execute(mock_request_context, mock_event_queue)  # type: ignore
+
+            # Assert
+            mock_updater.update_status.assert_called()
+            call_args_list = mock_updater.update_status.call_args_list
+            assert any(call[1].get("state") == TaskState.TASK_STATE_CANCELED for call in call_args_list)
+
+    async def test_execute_handles_generic_exception(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+        mock_task: Task,
+    ) -> None:
+        """Arrange: Create executor that raises generic exception
+        Act: Call execute method
+        Assert: Error is caught and task is marked as failed
+        """
+        # Arrange
+        mock_request_context.get_user_input = MagicMock(return_value="Hello")
+        mock_request_context.current_task = mock_task
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = MagicMock()
+
+        error_message = "Test error"
+        cast(Any, executor._agent).run = AsyncMock(side_effect=ValueError(error_message))
+        cast(Any, executor._agent).create_session = MagicMock()
+
+        with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+            mock_updater = MagicMock()
+            mock_updater.submit = AsyncMock()
+            mock_updater.start_work = AsyncMock()
+            mock_updater.update_status = AsyncMock()
+            mock_updater.new_agent_message = MagicMock(return_value="error_message_obj")
+            mock_updater_class.return_value = mock_updater
+
+            # Act
+            await executor.execute(mock_request_context, mock_event_queue)
+
+            # Assert
+            mock_updater.new_agent_message.assert_called_once()
+            args, _ = mock_updater.new_agent_message.call_args
+            parts = args[0]
+            assert len(parts) == 1
+            assert isinstance(parts[0], Part)
+            assert parts[0].text == error_message
+
+            call_args_list = mock_updater.update_status.call_args_list
+            assert any(
+                call[1].get("state") == TaskState.TASK_STATE_FAILED and call[1].get("message") == "error_message_obj"
+                for call in call_args_list
+            )
+
+    async def test_execute_processes_multiple_response_messages(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+        mock_task: Task,
+    ) -> None:
+        """Arrange: Create executor that returns multiple response messages
+        Act: Call execute method
+        Assert: All messages are processed through handle_events
+        """
+        # Arrange
+        mock_request_context.get_user_input = MagicMock(return_value="Hello")
+        mock_request_context.current_task = mock_task
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = MagicMock()
+
+        response_message1 = Message(role="assistant", contents=[Content.from_text(text="First")])
+        response_message2 = Message(role="assistant", contents=[Content.from_text(text="Second")])
+        response = MagicMock(spec=AgentResponse)
+        response.messages = [response_message1, response_message2]
+        cast(Any, executor._agent).run = AsyncMock(return_value=response)
+        cast(Any, executor._agent).create_session = MagicMock()
+
+        # Mock handle_events
+        cast(Any, executor).handle_events = AsyncMock()
+
+        with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+            mock_updater = MagicMock()
+            mock_updater.submit = AsyncMock()
+            mock_updater.start_work = AsyncMock()
+            mock_updater.complete = AsyncMock()
+            mock_updater.update_status = AsyncMock()
+            mock_updater_class.return_value = mock_updater
+
+            # Act
+            await executor.execute(mock_request_context, mock_event_queue)
+
+            # Assert
+            assert cast(Any, executor.handle_events).call_count == 2
+
+    async def test_execute_passes_query_to_run(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+        mock_task: Task,
+    ) -> None:
+        """Arrange: Create executor with request
+        Act: Call execute method
+        Assert: Query text is passed to run method with default stream and kwargs
+        """
+        # Arrange
+        query_text = "Hello agent"
+        mock_request_context.get_user_input = MagicMock(return_value=query_text)
+        mock_request_context.current_task = mock_task
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = MagicMock()
+
+        response_message = Message(role="assistant", contents=[Content.from_text(text="Response")])
+        response = MagicMock(spec=AgentResponse)
+        response.messages = [response_message]
+        cast(Any, executor._agent).run = AsyncMock(return_value=response)
+        cast(Any, executor._agent).create_session = MagicMock()
+
+        with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+            mock_updater = MagicMock()
+            mock_updater.submit = AsyncMock()
+            mock_updater.start_work = AsyncMock()
+            mock_updater.complete = AsyncMock()
+            mock_updater.update_status = AsyncMock()
+            mock_updater.new_agent_message = MagicMock(return_value="message_obj")
+            mock_updater_class.return_value = mock_updater
+
+            # Act
+            await executor.execute(mock_request_context, mock_event_queue)
+
+            # Assert
+            cast(Any, executor._agent.run).assert_called_once_with(
+                query_text, session=executor._agent.create_session(), stream=False
+            )
+
+    async def test_execute_with_stream_enabled(
+        self,
+        mock_agent: MagicMock,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+        mock_task: Task,
+    ) -> None:
+        """Arrange: Create executor with stream=True
+        Act: Call execute method
+        Assert: _run_stream is called and passes stream=True to run
+        """
+        # Arrange
+        executor = A2AExecutor(agent=mock_agent, stream=True)
+        query_text = "Hello agent"
+        mock_request_context.get_user_input = MagicMock(return_value=query_text)
+        mock_request_context.current_task = mock_task
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = MagicMock()
+
+        mock_response_stream = MagicMock()
+        mock_response_stream.with_transform_hook = MagicMock(return_value=mock_response_stream)
+        mock_response_stream.get_final_response = AsyncMock()
+        mock_agent.run = MagicMock(return_value=mock_response_stream)
+        mock_agent.create_session = MagicMock()
+
+        with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+            mock_updater = MagicMock()
+            mock_updater.submit = AsyncMock()
+            mock_updater.start_work = AsyncMock()
+            mock_updater.complete = AsyncMock()
+            mock_updater.update_status = AsyncMock()
+            mock_updater_class.return_value = mock_updater
+
+            # Act
+            await executor.execute(mock_request_context, mock_event_queue)
+
+            # Assert
+            mock_agent.run.assert_called_once_with(query_text, session=mock_agent.create_session(), stream=True)
+            mock_response_stream.with_transform_hook.assert_called_once()
+            mock_response_stream.get_final_response.assert_called_once()
+
+    async def test_execute_with_run_kwargs(
+        self,
+        mock_agent: MagicMock,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+        mock_task: Task,
+    ) -> None:
+        """Arrange: Create executor with run_kwargs
+        Act: Call execute method
+        Assert: run_kwargs are passed to run method
+        """
+        # Arrange
+        run_kwargs = {"temperature": 0.5, "max_tokens": 100}
+        executor = A2AExecutor(agent=mock_agent, run_kwargs=run_kwargs)
+        query_text = "Hello agent"
+        mock_request_context.get_user_input = MagicMock(return_value=query_text)
+        mock_request_context.current_task = mock_task
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = MagicMock()
+
+        response_message = Message(role="assistant", contents=[Content.from_text(text="Response")])
+        response = MagicMock(spec=AgentResponse)
+        response.messages = [response_message]
+        mock_agent.run = AsyncMock(return_value=response)
+        mock_agent.create_session = MagicMock()
+
+        with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+            mock_updater = MagicMock()
+            mock_updater.submit = AsyncMock()
+            mock_updater.start_work = AsyncMock()
+            mock_updater.complete = AsyncMock()
+            mock_updater.update_status = AsyncMock()
+            mock_updater_class.return_value = mock_updater
+
+            # Act
+            await executor.execute(mock_request_context, mock_event_queue)
+
+            # Assert
+            mock_agent.run.assert_called_once_with(
+                query_text, session=mock_agent.create_session(), stream=False, **run_kwargs
+            )
+
+
+class TestA2AExecutorHandleEvents:
+    """Tests for A2AExecutor.handle_events method."""
+
+    async def test_run_method_with_single_message(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test the private _run method with a single message (not a list)."""
+        # Arrange
+        query = "test query"
+        session = MagicMock()
+        response_message = Message(role="assistant", contents=[Content.from_text(text="Response")])
+        response = MagicMock(spec=AgentResponse)
+        response.messages = response_message  # Not a list
+        cast(Any, executor._agent).run = AsyncMock(return_value=response)
+        cast(Any, executor).handle_events = AsyncMock()
+
+        # Act
+        await executor._run(query, session, mock_updater)
+
+        # Assert
+        cast(Any, executor.handle_events).assert_called_once_with(response_message, mock_updater)
+
+    @fixture
+    def mock_updater(self) -> MagicMock:
+        """Create a mock execution context."""
+        updater = MagicMock()
+        updater.update_status = AsyncMock()
+        updater.new_agent_message = MagicMock(return_value="mock_message")
+        return updater
+
+    async def test_ignore_user_messages(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test that messages from USER role are ignored."""
+        # Arrange
+        message = Message(
+            contents=[Content.from_text(text="User input")],
+            role="user",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_not_called()
+
+    async def test_ignore_messages_with_no_contents(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test that messages with no contents are ignored."""
+        # Arrange
+        message = Message(
+            contents=[],
+            role="assistant",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_not_called()
+
+    async def test_handle_text_content(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test handling messages with text content."""
+        # Arrange
+        text = "Hello, this is a test message"
+        message = Message(
+            contents=[Content.from_text(text=text)],
+            role="assistant",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_called_once()
+        call_args = mock_updater.update_status.call_args
+        assert call_args.kwargs["state"] == TaskState.TASK_STATE_WORKING
+        assert mock_updater.new_agent_message.called
+
+    async def test_handle_multiple_text_contents(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test handling messages with multiple text contents."""
+        # Arrange
+        message = Message(
+            contents=[
+                Content.from_text(text="First message"),
+                Content.from_text(text="Second message"),
+            ],
+            role="assistant",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_called_once()
+        assert mock_updater.new_agent_message.called
+
+    async def test_handle_data_content(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test handling messages with data content."""
+        # Arrange
+        data = b"test file data"
+        message = Message(
+            contents=[Content.from_data(data=data, media_type="application/octet-stream")],
+            role="assistant",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_called_once()
+        call_args = mock_updater.update_status.call_args
+        assert call_args.kwargs["state"] == TaskState.TASK_STATE_WORKING
+
+    async def test_handle_uri_content(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test handling messages with URI content."""
+        # Arrange
+        uri = "https://example.com/file.pdf"
+        message = Message(
+            contents=[Content.from_uri(uri=uri, media_type="application/pdf")],
+            role="assistant",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_called_once()
+        call_args = mock_updater.update_status.call_args
+        assert call_args.kwargs["state"] == TaskState.TASK_STATE_WORKING
+
+    async def test_handle_mixed_content_types(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test handling messages with mixed content types."""
+        # Arrange
+        data = b"file data"
+
+        message = Message(
+            contents=[
+                Content.from_text(text="Processing file..."),
+                Content.from_data(data=data, media_type="application/octet-stream"),
+                Content.from_uri(uri="https://example.com/reference.pdf", media_type="application/pdf"),
+            ],
+            role="assistant",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_called_once()
+        call_args = mock_updater.update_status.call_args
+        assert call_args.kwargs["state"] == TaskState.TASK_STATE_WORKING
+
+    async def test_handle_with_additional_properties(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test handling messages with additional properties metadata."""
+        # Arrange
+        additional_props = {"custom_field": "custom_value", "priority": "high"}
+        message = Message(
+            contents=[Content.from_text(text="Test message")],
+            role="assistant",
+            additional_properties=additional_props,
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_called_once()
+        mock_updater.new_agent_message.assert_called_once()
+        call_args = mock_updater.new_agent_message.call_args
+        assert call_args.kwargs["metadata"] == additional_props
+
+    async def test_handle_with_no_additional_properties(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test handling messages without additional properties."""
+        # Arrange
+        message = Message(
+            contents=[Content.from_text(text="Test message")],
+            role="assistant",
+            additional_properties=None,
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_called_once()
+        mock_updater.new_agent_message.assert_called_once()
+        call_args = mock_updater.new_agent_message.call_args
+        assert call_args.kwargs["metadata"] == {}
+
+    async def test_parts_list_passed_to_new_agent_message(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test that parts list is correctly passed to new_agent_message."""
+        # Arrange
+        message = Message(
+            contents=[
+                Content.from_text(text="Message 1"),
+                Content.from_text(text="Message 2"),
+            ],
+            role="assistant",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.new_agent_message.assert_called_once()
+        call_kwargs = mock_updater.new_agent_message.call_args.kwargs
+        assert "parts" in call_kwargs
+        parts_list = call_kwargs["parts"]
+        assert len(parts_list) == 2
+
+    async def test_task_state_always_working(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test that task state is always set to working."""
+        # Arrange
+        message = Message(
+            contents=[Content.from_text(text="Any message")],
+            role="assistant",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        call_kwargs = mock_updater.update_status.call_args.kwargs
+        assert call_kwargs["state"] == TaskState.TASK_STATE_WORKING
+
+    async def test_handle_agent_response_update_no_streamed_set(
+        self, executor: A2AExecutor, mock_updater: MagicMock
+    ) -> None:
+        """Test handling AgentResponseUpdate (streaming) without a tracking set."""
+        # Arrange
+        update = AgentResponseUpdate(
+            contents=[Content.from_text(text="Streaming chunk")],
+            role="assistant",
+            message_id="msg-1",
+        )
+        mock_updater.add_artifact = AsyncMock()
+
+        # Act
+        await executor.handle_events(update, mock_updater)
+
+        # Assert
+        mock_updater.add_artifact.assert_called_once()
+        call_kwargs = mock_updater.add_artifact.call_args.kwargs
+        assert call_kwargs["artifact_id"] == "msg-1"
+        assert call_kwargs["append"] is None
+
+    async def test_handle_agent_response_update_first_time(
+        self, executor: A2AExecutor, mock_updater: MagicMock
+    ) -> None:
+        """Test handling AgentResponseUpdate (streaming) for the first time with a tracking set."""
+        # Arrange
+        update = AgentResponseUpdate(
+            contents=[Content.from_text(text="Streaming chunk")],
+            role="assistant",
+            message_id="msg-1",
+        )
+        mock_updater.add_artifact = AsyncMock()
+        streamed_artifact_ids: set[str] = set()
+
+        # Act
+        await executor.handle_events(update, mock_updater, streamed_artifact_ids=streamed_artifact_ids)
+
+        # Assert
+        mock_updater.add_artifact.assert_called_once()
+        call_kwargs = mock_updater.add_artifact.call_args.kwargs
+        assert call_kwargs["append"] is None
+        assert "msg-1" in streamed_artifact_ids
+
+    async def test_handle_agent_response_update_subsequent_time(
+        self, executor: A2AExecutor, mock_updater: MagicMock
+    ) -> None:
+        """Test handling AgentResponseUpdate (streaming) for subsequent times with a tracking set."""
+        # Arrange
+        update = AgentResponseUpdate(
+            contents=[Content.from_text(text="Next chunk")],
+            role="assistant",
+            message_id="msg-1",
+        )
+        mock_updater.add_artifact = AsyncMock()
+        streamed_artifact_ids = {"msg-1"}
+
+        # Act
+        await executor.handle_events(update, mock_updater, streamed_artifact_ids=streamed_artifact_ids)
+
+        # Assert
+        mock_updater.add_artifact.assert_called_once()
+        call_kwargs = mock_updater.add_artifact.call_args.kwargs
+        assert call_kwargs["append"] is True
+
+    async def test_handle_unsupported_content_type(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test handling messages with unsupported content types."""
+        # Arrange
+        message = Message(
+            contents=[Content(type=cast(Any, "unknown"), text="Some text")],  # type: ignore[arg-type]
+            role="assistant",
+        )
+
+        # Act
+        with patch("agent_framework_a2a._a2a_executor.logger") as mock_logger:
+            await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_logger.warning.assert_called_once()
+        mock_updater.update_status.assert_not_called()
+
+
+class TestA2AExecutorIntegration:
+    """Integration tests for A2AExecutor."""
+
+    async def test_full_execution_flow_with_responses(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+        mock_task: Task,
+    ) -> None:
+        """Arrange: Create executor with all mocked dependencies
+        Act: Execute full flow from request to completion
+        Assert: All components interact correctly
+        """
+        # Arrange
+        mock_request_context.get_user_input = MagicMock(return_value="Hello agent")
+        mock_request_context.current_task = mock_task
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = MagicMock()
+
+        response = MagicMock(spec=AgentResponse)
+        response_message = MagicMock(spec=Message)
+        response.messages = [response_message]
+        response_message.contents = [Content.from_text(text="Hello user")]
+        response_message.role = "assistant"
+        response_message.additional_properties = None
+
+        cast(Any, executor._agent).run = AsyncMock(return_value=response)
+        cast(Any, executor._agent).create_session = MagicMock()
+        cast(Any, executor).handle_events = AsyncMock()
+
+        with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+            mock_updater = MagicMock()
+            mock_updater.submit = AsyncMock()
+            mock_updater.start_work = AsyncMock()
+            mock_updater.complete = AsyncMock()
+            mock_updater.update_status = AsyncMock()
+            mock_updater_class.return_value = mock_updater
+
+            # Act
+            await executor.execute(mock_request_context, mock_event_queue)
+
+            # Assert
+            mock_updater.submit.assert_called_once()
+            mock_updater.start_work.assert_called_once()
+            cast(Any, executor.handle_events).assert_called_once()
+            mock_updater.complete.assert_called_once()

python/packages/a2a/tests/test_utils.py

@@ -0,0 +1,51 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import pytest
+
+from agent_framework_a2a._utils import get_uri_data
+
+
+def test_get_uri_data_valid() -> None:
+    """Test get_uri_data with valid data URIs."""
+    # Simple text/plain
+    uri = "data:text/plain;base64,SGVsbG8sIFdvcmxkIQ=="
+    assert get_uri_data(uri) == "SGVsbG8sIFdvcmxkIQ=="
+
+    # Image png
+    uri = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
+    assert get_uri_data(uri) == "iVBORw0KGgoAAAANSUhEUgfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
+
+    # Application octet-stream
+    uri = "data:application/octet-stream;base64,AQIDBA=="
+    assert get_uri_data(uri) == "AQIDBA=="
+
+    # Media type with parameters
+    uri = "data:text/plain;charset=utf-8;base64,SGVsbG8sIFdvcmxkIQ=="
+    assert get_uri_data(uri) == "SGVsbG8sIFdvcmxkIQ=="
+
+    # Media type with multiple parameters
+    uri = "data:text/plain;charset=utf-8;name=hello.txt;base64,SGVsbG8sIFdvcmxkIQ=="
+    assert get_uri_data(uri) == "SGVsbG8sIFdvcmxkIQ=="
+
+
+def test_get_uri_data_invalid_format() -> None:
+    """Test get_uri_data with invalid URI formats."""
+    invalid_uris = [
+        "not-a-uri",
+        "http://example.com",
+        "data:text/plain;SGVsbG8sIFdvcmxkIQ==",  # Missing base64 marker
+        "data:base64,SGVsbG8sIFdvcmxkIQ==",  # Missing media type
+        "data:text/plain;foo;base64,SGVsbG8sIFdvcmxkIQ==",  # Parameter without value
+        "data:text/plain;base64;base64,SGVsbG8sIFdvcmxkIQ==",  # base64 used as a parameter name
+        "data:text/plain;base64,SGVsbG8sIFdvcmxkIQ== extra",
+        "data:text/plain;base64,SGVsbG8sIFdvcmxkIQ==\n",
+    ]
+    for uri in invalid_uris:
+        with pytest.raises(ValueError, match="Invalid data URI format"):
+            get_uri_data(uri)
+
+
+def test_get_uri_data_empty() -> None:
+    """Test get_uri_data with empty string."""
+    with pytest.raises(ValueError, match="Invalid data URI format"):
+        get_uri_data("")