39
39
from a2a .types import AgentCapabilities
40
40
from a2a .types import AgentCard
41
41
from a2a .types import AgentSkill
42
+ from a2a .types import Artifact
42
43
from a2a .types import Message as A2AMessage
43
44
from a2a .types import SendMessageSuccessResponse
44
45
from a2a .types import Task as A2ATask
46
+ from a2a .types import TaskArtifactUpdateEvent
47
+ from a2a .types import TaskStatusUpdateEvent
45
48
from google .adk .agents .invocation_context import InvocationContext
46
49
from google .adk .agents .remote_a2a_agent import A2A_METADATA_PREFIX
47
50
from google .adk .agents .remote_a2a_agent import AgentCardResolutionError
@@ -60,6 +63,9 @@ class DummyTypes:
60
63
A2AMessage = DummyTypes ()
61
64
SendMessageSuccessResponse = DummyTypes ()
62
65
A2ATask = DummyTypes ()
66
+ TaskStatusUpdateEvent = DummyTypes ()
67
+ Artifact = DummyTypes ()
68
+ TaskArtifactUpdateEvent = DummyTypes ()
63
69
InvocationContext = DummyTypes ()
64
70
RemoteA2aAgent = DummyTypes ()
65
71
AgentCardResolutionError = Exception
@@ -685,8 +691,8 @@ async def test_handle_a2a_response_success_with_message(self):
685
691
assert A2A_METADATA_PREFIX + "context_id" in result .custom_metadata
686
692
687
693
@pytest .mark .asyncio
688
- async def test_handle_a2a_response_success_with_task (self ):
689
- """Test successful A2A response handling with task."""
694
+ async def test_handle_a2a_response_with_task_and_no_update (self ):
695
+ """Test successful A2A response handling with task and no update ."""
690
696
mock_a2a_task = Mock (spec = A2ATask )
691
697
mock_a2a_task .id = "task-123"
692
698
mock_a2a_task .context_id = "context-123"
@@ -718,6 +724,116 @@ async def test_handle_a2a_response_success_with_task(self):
718
724
assert A2A_METADATA_PREFIX + "task_id" in result .custom_metadata
719
725
assert A2A_METADATA_PREFIX + "context_id" in result .custom_metadata
720
726
727
+ @pytest .mark .asyncio
728
+ async def test_handle_a2a_response_with_task_status_update_with_message (self ):
729
+ """Test handling of a task status update with a message."""
730
+ mock_a2a_task = Mock (spec = A2ATask )
731
+ mock_a2a_task .id = "task-123"
732
+ mock_a2a_task .context_id = "context-123"
733
+
734
+ mock_a2a_message = Mock (spec = A2AMessage )
735
+ mock_update = Mock (spec = TaskStatusUpdateEvent )
736
+ mock_update .message = mock_a2a_message
737
+ mock_update .status = "COMPLETED"
738
+
739
+ # Create a proper Event mock that can handle custom_metadata
740
+ mock_event = Event (
741
+ author = self .agent .name ,
742
+ invocation_id = self .mock_context .invocation_id ,
743
+ branch = self .mock_context .branch ,
744
+ )
745
+
746
+ with patch (
747
+ "google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
748
+ ) as mock_convert :
749
+ mock_convert .return_value = mock_event
750
+
751
+ result = await self .agent ._handle_a2a_response (
752
+ (mock_a2a_task , mock_update ), self .mock_context
753
+ )
754
+
755
+ assert result == mock_event
756
+ mock_convert .assert_called_once_with (
757
+ mock_a2a_message ,
758
+ self .agent .name ,
759
+ self .mock_context ,
760
+ )
761
+ # Check that metadata was added
762
+ assert result .custom_metadata is not None
763
+ assert A2A_METADATA_PREFIX + "task_id" in result .custom_metadata
764
+ assert A2A_METADATA_PREFIX + "context_id" in result .custom_metadata
765
+
766
+ @pytest .mark .asyncio
767
+ async def test_handle_a2a_response_with_task_status_update_no_message (self ):
768
+ """Test handling of a task status update with no message."""
769
+ mock_a2a_task = Mock (spec = A2ATask )
770
+ mock_a2a_task .id = "task-123"
771
+
772
+ mock_update = Mock (spec = TaskStatusUpdateEvent )
773
+ mock_update .message = None
774
+ mock_update .status = "COMPLETED"
775
+
776
+ result = await self .agent ._handle_a2a_response (
777
+ (mock_a2a_task , mock_update ), self .mock_context
778
+ )
779
+
780
+ assert result is None
781
+
782
+ @pytest .mark .asyncio
783
+ async def test_handle_a2a_response_with_artifact_update (self ):
784
+ """Test successful A2A response handling with artifact update."""
785
+ mock_a2a_task = Mock (spec = A2ATask )
786
+ mock_a2a_task .id = "task-123"
787
+ mock_a2a_task .context_id = "context-123"
788
+
789
+ mock_artifact = Mock (spec = Artifact )
790
+ mock_update = Mock (spec = TaskArtifactUpdateEvent )
791
+ mock_update .artifact = mock_artifact
792
+ mock_update .append = False
793
+ mock_update .last_chunk = True
794
+
795
+ # Create a proper Event mock that can handle custom_metadata
796
+ mock_event = Event (
797
+ author = self .agent .name ,
798
+ invocation_id = self .mock_context .invocation_id ,
799
+ branch = self .mock_context .branch ,
800
+ )
801
+
802
+ with patch (
803
+ "google.adk.agents.remote_a2a_agent.convert_a2a_task_to_event"
804
+ ) as mock_convert :
805
+ mock_convert .return_value = mock_event
806
+
807
+ result = await self .agent ._handle_a2a_response (
808
+ (mock_a2a_task , mock_update ), self .mock_context
809
+ )
810
+
811
+ assert result == mock_event
812
+ mock_convert .assert_called_once_with (
813
+ mock_a2a_task , self .agent .name , self .mock_context
814
+ )
815
+ # Check that metadata was added
816
+ assert result .custom_metadata is not None
817
+ assert A2A_METADATA_PREFIX + "task_id" in result .custom_metadata
818
+ assert A2A_METADATA_PREFIX + "context_id" in result .custom_metadata
819
+
820
+ @pytest .mark .asyncio
821
+ async def test_handle_a2a_response_with_partial_artifact_update (self ):
822
+ """Test that partial artifact updates are ignored."""
823
+ mock_a2a_task = Mock (spec = A2ATask )
824
+ mock_a2a_task .id = "task-123"
825
+
826
+ mock_update = Mock (spec = TaskArtifactUpdateEvent )
827
+ mock_update .artifact = Mock (spec = Artifact )
828
+ mock_update .append = True
829
+ mock_update .last_chunk = False
830
+
831
+ result = await self .agent ._handle_a2a_response (
832
+ (mock_a2a_task , mock_update ), self .mock_context
833
+ )
834
+
835
+ assert result is None
836
+
721
837
722
838
class TestRemoteA2aAgentMessageHandlingFromFactory :
723
839
"""Test message handling functionality."""
@@ -865,8 +981,8 @@ async def test_handle_a2a_response_success_with_message(self):
865
981
assert A2A_METADATA_PREFIX + "context_id" in result .custom_metadata
866
982
867
983
@pytest .mark .asyncio
868
- async def test_handle_a2a_response_success_with_task (self ):
869
- """Test successful A2A response handling with task."""
984
+ async def test_handle_a2a_response_with_task_and_no_update (self ):
985
+ """Test successful A2A response handling with task and no update ."""
870
986
mock_a2a_task = Mock (spec = A2ATask )
871
987
mock_a2a_task .id = "task-123"
872
988
mock_a2a_task .context_id = "context-123"
@@ -896,6 +1012,116 @@ async def test_handle_a2a_response_success_with_task(self):
896
1012
assert A2A_METADATA_PREFIX + "task_id" in result .custom_metadata
897
1013
assert A2A_METADATA_PREFIX + "context_id" in result .custom_metadata
898
1014
1015
+ @pytest .mark .asyncio
1016
+ async def test_handle_a2a_response_with_task_status_update_with_message (self ):
1017
+ """Test handling of a task status update with a message."""
1018
+ mock_a2a_task = Mock (spec = A2ATask )
1019
+ mock_a2a_task .id = "task-123"
1020
+ mock_a2a_task .context_id = "context-123"
1021
+
1022
+ mock_a2a_message = Mock (spec = A2AMessage )
1023
+ mock_update = Mock (spec = TaskStatusUpdateEvent )
1024
+ mock_update .message = mock_a2a_message
1025
+ mock_update .status = "COMPLETED"
1026
+
1027
+ # Create a proper Event mock that can handle custom_metadata
1028
+ mock_event = Event (
1029
+ author = self .agent .name ,
1030
+ invocation_id = self .mock_context .invocation_id ,
1031
+ branch = self .mock_context .branch ,
1032
+ )
1033
+
1034
+ with patch (
1035
+ "google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
1036
+ ) as mock_convert :
1037
+ mock_convert .return_value = mock_event
1038
+
1039
+ result = await self .agent ._handle_a2a_response (
1040
+ (mock_a2a_task , mock_update ), self .mock_context
1041
+ )
1042
+
1043
+ assert result == mock_event
1044
+ mock_convert .assert_called_once_with (
1045
+ mock_a2a_message ,
1046
+ self .agent .name ,
1047
+ self .mock_context ,
1048
+ )
1049
+ # Check that metadata was added
1050
+ assert result .custom_metadata is not None
1051
+ assert A2A_METADATA_PREFIX + "task_id" in result .custom_metadata
1052
+ assert A2A_METADATA_PREFIX + "context_id" in result .custom_metadata
1053
+
1054
+ @pytest .mark .asyncio
1055
+ async def test_handle_a2a_response_with_task_status_update_no_message (self ):
1056
+ """Test handling of a task status update with no message."""
1057
+ mock_a2a_task = Mock (spec = A2ATask )
1058
+ mock_a2a_task .id = "task-123"
1059
+
1060
+ mock_update = Mock (spec = TaskStatusUpdateEvent )
1061
+ mock_update .message = None
1062
+ mock_update .status = "COMPLETED"
1063
+
1064
+ result = await self .agent ._handle_a2a_response (
1065
+ (mock_a2a_task , mock_update ), self .mock_context
1066
+ )
1067
+
1068
+ assert result is None
1069
+
1070
+ @pytest .mark .asyncio
1071
+ async def test_handle_a2a_response_with_artifact_update (self ):
1072
+ """Test successful A2A response handling with artifact update."""
1073
+ mock_a2a_task = Mock (spec = A2ATask )
1074
+ mock_a2a_task .id = "task-123"
1075
+ mock_a2a_task .context_id = "context-123"
1076
+
1077
+ mock_artifact = Mock (spec = Artifact )
1078
+ mock_update = Mock (spec = TaskArtifactUpdateEvent )
1079
+ mock_update .artifact = mock_artifact
1080
+ mock_update .append = False
1081
+ mock_update .last_chunk = True
1082
+
1083
+ # Create a proper Event mock that can handle custom_metadata
1084
+ mock_event = Event (
1085
+ author = self .agent .name ,
1086
+ invocation_id = self .mock_context .invocation_id ,
1087
+ branch = self .mock_context .branch ,
1088
+ )
1089
+
1090
+ with patch (
1091
+ "google.adk.agents.remote_a2a_agent.convert_a2a_task_to_event"
1092
+ ) as mock_convert :
1093
+ mock_convert .return_value = mock_event
1094
+
1095
+ result = await self .agent ._handle_a2a_response (
1096
+ (mock_a2a_task , mock_update ), self .mock_context
1097
+ )
1098
+
1099
+ assert result == mock_event
1100
+ mock_convert .assert_called_once_with (
1101
+ mock_a2a_task , self .agent .name , self .mock_context
1102
+ )
1103
+ # Check that metadata was added
1104
+ assert result .custom_metadata is not None
1105
+ assert A2A_METADATA_PREFIX + "task_id" in result .custom_metadata
1106
+ assert A2A_METADATA_PREFIX + "context_id" in result .custom_metadata
1107
+
1108
+ @pytest .mark .asyncio
1109
+ async def test_handle_a2a_response_with_partial_artifact_update (self ):
1110
+ """Test that partial artifact updates are ignored."""
1111
+ mock_a2a_task = Mock (spec = A2ATask )
1112
+ mock_a2a_task .id = "task-123"
1113
+
1114
+ mock_update = Mock (spec = TaskArtifactUpdateEvent )
1115
+ mock_update .artifact = Mock (spec = Artifact )
1116
+ mock_update .append = True
1117
+ mock_update .last_chunk = False
1118
+
1119
+ result = await self .agent ._handle_a2a_response (
1120
+ (mock_a2a_task , mock_update ), self .mock_context
1121
+ )
1122
+
1123
+ assert result is None
1124
+
899
1125
900
1126
class TestRemoteA2aAgentExecution :
901
1127
"""Test agent execution functionality."""
@@ -1019,6 +1245,7 @@ async def test_run_async_impl_successful_request(self):
1019
1245
# Add model_dump to mock_response for metadata
1020
1246
mock_response .model_dump .return_value = {"test" : "response" }
1021
1247
1248
+ # Execute
1022
1249
events = []
1023
1250
async for event in self .agent ._run_async_impl (
1024
1251
self .mock_context
@@ -1211,6 +1438,7 @@ async def test_run_async_impl_successful_request(self):
1211
1438
"test" : "response"
1212
1439
}
1213
1440
1441
+ # Execute
1214
1442
events = []
1215
1443
async for event in self .agent ._run_async_impl (
1216
1444
self .mock_context
0 commit comments