16
16
from a2a .server .events import QueueManager
17
17
from a2a .server .events .event_queue import EventQueue
18
18
from a2a .server .request_handlers import DefaultRequestHandler , JSONRPCHandler
19
- from a2a .server .tasks import TaskStore , InMemoryPushNotificationConfigStore , BasePushNotificationSender , PushNotificationConfigStore , PushNotificationSender
19
+ from a2a .server .tasks import (
20
+ TaskStore ,
21
+ InMemoryPushNotificationConfigStore ,
22
+ BasePushNotificationSender ,
23
+ PushNotificationConfigStore ,
24
+ PushNotificationSender ,
25
+ )
20
26
from a2a .types import (
21
27
AgentCapabilities ,
22
28
AgentCard ,
@@ -436,8 +442,10 @@ async def streaming_coro():
436
442
async def test_set_push_notification_success (self ) -> None :
437
443
mock_agent_executor = AsyncMock (spec = AgentExecutor )
438
444
mock_task_store = AsyncMock (spec = TaskStore )
439
- mock_push_notification_store = AsyncMock (spec = PushNotificationConfigStore )
440
-
445
+ mock_push_notification_store = AsyncMock (
446
+ spec = PushNotificationConfigStore
447
+ )
448
+
441
449
request_handler = DefaultRequestHandler (
442
450
mock_agent_executor ,
443
451
mock_task_store ,
@@ -471,10 +479,12 @@ async def test_set_push_notification_success(self) -> None:
471
479
472
480
async def test_get_push_notification_success (self ) -> None :
473
481
mock_agent_executor = AsyncMock (spec = AgentExecutor )
474
- mock_task_store = AsyncMock (spec = TaskStore )
482
+ mock_task_store = AsyncMock (spec = TaskStore )
475
483
push_notification_store = InMemoryPushNotificationConfigStore ()
476
484
request_handler = DefaultRequestHandler (
477
- mock_agent_executor , mock_task_store , push_config_store = push_notification_store
485
+ mock_agent_executor ,
486
+ mock_task_store ,
487
+ push_config_store = push_notification_store ,
478
488
)
479
489
self .mock_agent_card .capabilities = AgentCapabilities (
480
490
streaming = True , pushNotifications = True
@@ -516,9 +526,14 @@ async def test_on_message_stream_new_message_send_push_notification_success(
516
526
mock_task_store = AsyncMock (spec = TaskStore )
517
527
mock_httpx_client = AsyncMock (spec = httpx .AsyncClient )
518
528
push_notification_store = InMemoryPushNotificationConfigStore ()
519
- push_notification_sender = BasePushNotificationSender (mock_httpx_client , push_notification_store )
529
+ push_notification_sender = BasePushNotificationSender (
530
+ mock_httpx_client , push_notification_store
531
+ )
520
532
request_handler = DefaultRequestHandler (
521
- mock_agent_executor , mock_task_store , push_config_store = push_notification_store , push_sender = push_notification_sender
533
+ mock_agent_executor ,
534
+ mock_task_store ,
535
+ push_config_store = push_notification_store ,
536
+ push_sender = push_notification_sender ,
522
537
)
523
538
self .mock_agent_card .capabilities = AgentCapabilities (
524
539
streaming = True , pushNotifications = True
@@ -585,7 +600,7 @@ async def streaming_coro():
585
600
'kind' : 'task' ,
586
601
'status' : {'state' : 'submitted' },
587
602
},
588
- headers = None
603
+ headers = None ,
589
604
),
590
605
call (
591
606
'http://example.com' ,
@@ -606,7 +621,7 @@ async def streaming_coro():
606
621
'kind' : 'task' ,
607
622
'status' : {'state' : 'submitted' },
608
623
},
609
- headers = None
624
+ headers = None ,
610
625
),
611
626
call (
612
627
'http://example.com' ,
@@ -627,7 +642,7 @@ async def streaming_coro():
627
642
'kind' : 'task' ,
628
643
'status' : {'state' : 'completed' },
629
644
},
630
- headers = None
645
+ headers = None ,
631
646
),
632
647
]
633
648
mock_httpx_client .post .assert_has_calls (calls )
@@ -727,7 +742,7 @@ async def test_streaming_not_supported_error(
727
742
pass
728
743
729
744
self .assertEqual (
730
- str (context .exception .error .message ), # type: ignore
745
+ str (context .exception .error .message ), # type: ignore
731
746
'Streaming is not supported by the agent' ,
732
747
)
733
748
@@ -761,7 +776,7 @@ async def test_push_notifications_not_supported_error(self) -> None:
761
776
await handler .set_push_notification_config (request )
762
777
763
778
self .assertEqual (
764
- str (context .exception .error .message ), # type: ignore
779
+ str (context .exception .error .message ), # type: ignore
765
780
'Push notifications are not supported by the agent' ,
766
781
)
767
782
@@ -960,7 +975,7 @@ async def consume_raises_error(*args, **kwargs) -> NoReturn:
960
975
961
976
# Assert
962
977
self .assertIsInstance (response .root , JSONRPCErrorResponse )
963
- self .assertEqual (response .root .error , UnsupportedOperationError ()) # type: ignore
978
+ self .assertEqual (response .root .error , UnsupportedOperationError ()) # type: ignore
964
979
965
980
async def test_on_message_send_task_id_mismatch (self ) -> None :
966
981
mock_agent_executor = AsyncMock (spec = AgentExecutor )
@@ -1031,37 +1046,54 @@ async def test_on_get_push_notification(self) -> None:
1031
1046
1032
1047
mock_task = Task (** MINIMAL_TASK )
1033
1048
mock_task_store .get .return_value = mock_task
1034
-
1035
1049
1036
1050
# Create request handler without a push notifier
1037
1051
request_handler = AsyncMock (spec = DefaultRequestHandler )
1038
- task_push_config = TaskPushNotificationConfig (taskId = mock_task .id , pushNotificationConfig = PushNotificationConfig (id = "config1" , url = 'http://example.com' ))
1039
- request_handler .on_get_task_push_notification_config .return_value = task_push_config
1052
+ task_push_config = TaskPushNotificationConfig (
1053
+ taskId = mock_task .id ,
1054
+ pushNotificationConfig = PushNotificationConfig (
1055
+ id = 'config1' , url = 'http://example.com'
1056
+ ),
1057
+ )
1058
+ request_handler .on_get_task_push_notification_config .return_value = (
1059
+ task_push_config
1060
+ )
1040
1061
1041
1062
self .mock_agent_card .capabilities = AgentCapabilities (
1042
1063
pushNotifications = True
1043
1064
)
1044
1065
handler = JSONRPCHandler (self .mock_agent_card , request_handler )
1045
1066
list_request = GetTaskPushNotificationConfigRequest (
1046
- id = '1' , params = GetTaskPushNotificationConfigParams (id = mock_task .id , pushNotificationConfigId = "config1" )
1067
+ id = '1' ,
1068
+ params = GetTaskPushNotificationConfigParams (
1069
+ id = mock_task .id , pushNotificationConfigId = 'config1'
1070
+ ),
1047
1071
)
1048
1072
response = await handler .get_push_notification_config (list_request )
1049
1073
# Assert
1050
- self .assertIsInstance (response .root , GetTaskPushNotificationConfigSuccessResponse )
1051
- self .assertEqual (response .root .result , task_push_config ) # type: ignore
1074
+ self .assertIsInstance (
1075
+ response .root , GetTaskPushNotificationConfigSuccessResponse
1076
+ )
1077
+ self .assertEqual (response .root .result , task_push_config ) # type: ignore
1052
1078
1053
1079
async def test_on_list_push_notification (self ) -> None :
1054
1080
"""Test list_push_notification_config handling"""
1055
1081
mock_task_store = AsyncMock (spec = TaskStore )
1056
1082
1057
1083
mock_task = Task (** MINIMAL_TASK )
1058
1084
mock_task_store .get .return_value = mock_task
1059
-
1060
1085
1061
1086
# Create request handler without a push notifier
1062
1087
request_handler = AsyncMock (spec = DefaultRequestHandler )
1063
- task_push_config = TaskPushNotificationConfig (taskId = mock_task .id , pushNotificationConfig = PushNotificationConfig (url = 'http://example.com' ))
1064
- request_handler .on_list_task_push_notification_config .return_value = [task_push_config ]
1088
+ task_push_config = TaskPushNotificationConfig (
1089
+ taskId = mock_task .id ,
1090
+ pushNotificationConfig = PushNotificationConfig (
1091
+ url = 'http://example.com'
1092
+ ),
1093
+ )
1094
+ request_handler .on_list_task_push_notification_config .return_value = [
1095
+ task_push_config
1096
+ ]
1065
1097
1066
1098
self .mock_agent_card .capabilities = AgentCapabilities (
1067
1099
pushNotifications = True
@@ -1072,23 +1104,30 @@ async def test_on_list_push_notification(self) -> None:
1072
1104
)
1073
1105
response = await handler .list_push_notification_config (list_request )
1074
1106
# Assert
1075
- self .assertIsInstance (response .root , ListTaskPushNotificationConfigSuccessResponse )
1076
- self .assertEqual (response .root .result , [task_push_config ]) # type: ignore
1107
+ self .assertIsInstance (
1108
+ response .root , ListTaskPushNotificationConfigSuccessResponse
1109
+ )
1110
+ self .assertEqual (response .root .result , [task_push_config ]) # type: ignore
1077
1111
1078
1112
async def test_on_list_push_notification_error (self ) -> None :
1079
1113
"""Test list_push_notification_config handling"""
1080
1114
mock_task_store = AsyncMock (spec = TaskStore )
1081
1115
1082
1116
mock_task = Task (** MINIMAL_TASK )
1083
1117
mock_task_store .get .return_value = mock_task
1084
-
1085
1118
1086
1119
# Create request handler without a push notifier
1087
1120
request_handler = AsyncMock (spec = DefaultRequestHandler )
1088
- task_push_config = TaskPushNotificationConfig (taskId = mock_task .id , pushNotificationConfig = PushNotificationConfig (url = 'http://example.com' ))
1121
+ task_push_config = TaskPushNotificationConfig (
1122
+ taskId = mock_task .id ,
1123
+ pushNotificationConfig = PushNotificationConfig (
1124
+ url = 'http://example.com'
1125
+ ),
1126
+ )
1089
1127
# throw server error
1090
- request_handler .on_list_task_push_notification_config .side_effect = ServerError (InternalError ())
1091
-
1128
+ request_handler .on_list_task_push_notification_config .side_effect = (
1129
+ ServerError (InternalError ())
1130
+ )
1092
1131
1093
1132
self .mock_agent_card .capabilities = AgentCapabilities (
1094
1133
pushNotifications = True
@@ -1100,45 +1139,55 @@ async def test_on_list_push_notification_error(self) -> None:
1100
1139
response = await handler .list_push_notification_config (list_request )
1101
1140
# Assert
1102
1141
self .assertIsInstance (response .root , JSONRPCErrorResponse )
1103
- self .assertEqual (response .root .error , InternalError ()) # type: ignore
1104
-
1142
+ self .assertEqual (response .root .error , InternalError ()) # type: ignore
1143
+
1105
1144
async def test_on_delete_push_notification (self ) -> None :
1106
1145
"""Test delete_push_notification_config handling"""
1107
1146
1108
1147
# Create request handler without a push notifier
1109
- request_handler = AsyncMock (spec = DefaultRequestHandler )
1110
- request_handler .on_delete_task_push_notification_config .return_value = None
1148
+ request_handler = AsyncMock (spec = DefaultRequestHandler )
1149
+ request_handler .on_delete_task_push_notification_config .return_value = (
1150
+ None
1151
+ )
1111
1152
1112
1153
self .mock_agent_card .capabilities = AgentCapabilities (
1113
1154
pushNotifications = True
1114
1155
)
1115
1156
handler = JSONRPCHandler (self .mock_agent_card , request_handler )
1116
1157
delete_request = DeleteTaskPushNotificationConfigRequest (
1117
- id = '1' , params = DeleteTaskPushNotificationConfigParams (id = "task1" , pushNotificationConfigId = "config1" )
1158
+ id = '1' ,
1159
+ params = DeleteTaskPushNotificationConfigParams (
1160
+ id = 'task1' , pushNotificationConfigId = 'config1'
1161
+ ),
1118
1162
)
1119
1163
response = await handler .delete_push_notification_config (delete_request )
1120
1164
# Assert
1121
- self .assertIsInstance (response .root , DeleteTaskPushNotificationConfigSuccessResponse )
1122
- self .assertEqual (response .root .result , None ) # type: ignore
1165
+ self .assertIsInstance (
1166
+ response .root , DeleteTaskPushNotificationConfigSuccessResponse
1167
+ )
1168
+ self .assertEqual (response .root .result , None ) # type: ignore
1123
1169
1124
1170
async def test_on_delete_push_notification_error (self ) -> None :
1125
1171
"""Test delete_push_notification_config error handling"""
1126
-
1127
1172
1128
1173
# Create request handler without a push notifier
1129
1174
request_handler = AsyncMock (spec = DefaultRequestHandler )
1130
1175
# throw server error
1131
- request_handler .on_delete_task_push_notification_config .side_effect = ServerError (UnsupportedOperationError ())
1132
-
1176
+ request_handler .on_delete_task_push_notification_config .side_effect = (
1177
+ ServerError (UnsupportedOperationError ())
1178
+ )
1133
1179
1134
1180
self .mock_agent_card .capabilities = AgentCapabilities (
1135
1181
pushNotifications = True
1136
1182
)
1137
1183
handler = JSONRPCHandler (self .mock_agent_card , request_handler )
1138
1184
delete_request = DeleteTaskPushNotificationConfigRequest (
1139
- id = '1' , params = DeleteTaskPushNotificationConfigParams (id = "task1" , pushNotificationConfigId = "config1" )
1185
+ id = '1' ,
1186
+ params = DeleteTaskPushNotificationConfigParams (
1187
+ id = 'task1' , pushNotificationConfigId = 'config1'
1188
+ ),
1140
1189
)
1141
1190
response = await handler .delete_push_notification_config (delete_request )
1142
1191
# Assert
1143
1192
self .assertIsInstance (response .root , JSONRPCErrorResponse )
1144
- self .assertEqual (response .root .error , UnsupportedOperationError ()) # type: ignore
1193
+ self .assertEqual (response .root .error , UnsupportedOperationError ()) # type: ignore
0 commit comments