Skip to content

Commit e7d6bb1

Browse files
amitami2amitami2
committed
added the unit tests
Signed-off-by: amitami2 <amitami2@cisco.com>
1 parent 98dda2f commit e7d6bb1

1 file changed

Lines changed: 41 additions & 199 deletions

File tree

‎a2a-websocket/src/server.rs‎

Lines changed: 41 additions & 199 deletions
Original file line numberDiff line numberDiff line change
@@ -671,11 +671,27 @@ mod tests {
671671
))
672672
}
673673

674-
struct StubHandler;
674+
#[derive(Default)]
675+
struct StubHandler {
676+
send_message_error: Option<A2AError>,
677+
streaming_pending: bool,
678+
}
675679

676-
struct FatalHandler;
680+
impl StubHandler {
681+
fn fatal_send_message() -> Self {
682+
Self {
683+
send_message_error: Some(A2AError::new(error_code::PARSE_ERROR, "fatal parse")),
684+
streaming_pending: false,
685+
}
686+
}
677687

678-
struct PendingStreamHandler;
688+
fn pending_stream() -> Self {
689+
Self {
690+
send_message_error: None,
691+
streaming_pending: true,
692+
}
693+
}
694+
}
679695

680696
fn sample_task(id: &str) -> Task {
681697
Task {
@@ -741,6 +757,10 @@ mod tests {
741757
_params: &ServiceParams,
742758
_req: SendMessageRequest,
743759
) -> Result<SendMessageResponse, A2AError> {
760+
if let Some(error) = &self.send_message_error {
761+
return Err(error.clone());
762+
}
763+
744764
Ok(SendMessageResponse::Task(sample_task("send")))
745765
}
746766

@@ -749,6 +769,10 @@ mod tests {
749769
_params: &ServiceParams,
750770
_req: SendMessageRequest,
751771
) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError> {
772+
if self.streaming_pending {
773+
return Ok(Box::pin(futures::stream::pending()));
774+
}
775+
752776
Ok(Box::pin(futures::stream::iter(vec![Ok(
753777
StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
754778
task_id: "stream".into(),
@@ -859,188 +883,6 @@ mod tests {
859883
}
860884
}
861885

862-
#[async_trait]
863-
impl RequestHandler for FatalHandler {
864-
async fn send_message(
865-
&self,
866-
_params: &ServiceParams,
867-
_req: SendMessageRequest,
868-
) -> Result<SendMessageResponse, A2AError> {
869-
Err(A2AError::new(error_code::PARSE_ERROR, "fatal parse"))
870-
}
871-
872-
async fn send_streaming_message(
873-
&self,
874-
_params: &ServiceParams,
875-
_req: SendMessageRequest,
876-
) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError> {
877-
unreachable!()
878-
}
879-
880-
async fn get_task(
881-
&self,
882-
_params: &ServiceParams,
883-
_req: GetTaskRequest,
884-
) -> Result<Task, A2AError> {
885-
unreachable!()
886-
}
887-
888-
async fn list_tasks(
889-
&self,
890-
_params: &ServiceParams,
891-
_req: ListTasksRequest,
892-
) -> Result<ListTasksResponse, A2AError> {
893-
unreachable!()
894-
}
895-
896-
async fn cancel_task(
897-
&self,
898-
_params: &ServiceParams,
899-
_req: CancelTaskRequest,
900-
) -> Result<Task, A2AError> {
901-
unreachable!()
902-
}
903-
904-
async fn subscribe_to_task(
905-
&self,
906-
_params: &ServiceParams,
907-
_req: SubscribeToTaskRequest,
908-
) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError> {
909-
unreachable!()
910-
}
911-
912-
async fn create_push_config(
913-
&self,
914-
_params: &ServiceParams,
915-
_req: CreateTaskPushNotificationConfigRequest,
916-
) -> Result<TaskPushNotificationConfig, A2AError> {
917-
unreachable!()
918-
}
919-
920-
async fn get_push_config(
921-
&self,
922-
_params: &ServiceParams,
923-
_req: GetTaskPushNotificationConfigRequest,
924-
) -> Result<TaskPushNotificationConfig, A2AError> {
925-
unreachable!()
926-
}
927-
928-
async fn list_push_configs(
929-
&self,
930-
_params: &ServiceParams,
931-
_req: ListTaskPushNotificationConfigsRequest,
932-
) -> Result<ListTaskPushNotificationConfigsResponse, A2AError> {
933-
unreachable!()
934-
}
935-
936-
async fn delete_push_config(
937-
&self,
938-
_params: &ServiceParams,
939-
_req: DeleteTaskPushNotificationConfigRequest,
940-
) -> Result<(), A2AError> {
941-
unreachable!()
942-
}
943-
944-
async fn get_extended_agent_card(
945-
&self,
946-
_params: &ServiceParams,
947-
_req: GetExtendedAgentCardRequest,
948-
) -> Result<AgentCard, A2AError> {
949-
unreachable!()
950-
}
951-
}
952-
953-
#[async_trait]
954-
impl RequestHandler for PendingStreamHandler {
955-
async fn send_message(
956-
&self,
957-
_params: &ServiceParams,
958-
_req: SendMessageRequest,
959-
) -> Result<SendMessageResponse, A2AError> {
960-
unreachable!()
961-
}
962-
963-
async fn send_streaming_message(
964-
&self,
965-
_params: &ServiceParams,
966-
_req: SendMessageRequest,
967-
) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError> {
968-
Ok(Box::pin(futures::stream::pending()))
969-
}
970-
971-
async fn get_task(
972-
&self,
973-
_params: &ServiceParams,
974-
_req: GetTaskRequest,
975-
) -> Result<Task, A2AError> {
976-
unreachable!()
977-
}
978-
979-
async fn list_tasks(
980-
&self,
981-
_params: &ServiceParams,
982-
_req: ListTasksRequest,
983-
) -> Result<ListTasksResponse, A2AError> {
984-
unreachable!()
985-
}
986-
987-
async fn cancel_task(
988-
&self,
989-
_params: &ServiceParams,
990-
_req: CancelTaskRequest,
991-
) -> Result<Task, A2AError> {
992-
unreachable!()
993-
}
994-
995-
async fn subscribe_to_task(
996-
&self,
997-
_params: &ServiceParams,
998-
_req: SubscribeToTaskRequest,
999-
) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError> {
1000-
unreachable!()
1001-
}
1002-
1003-
async fn create_push_config(
1004-
&self,
1005-
_params: &ServiceParams,
1006-
_req: CreateTaskPushNotificationConfigRequest,
1007-
) -> Result<TaskPushNotificationConfig, A2AError> {
1008-
unreachable!()
1009-
}
1010-
1011-
async fn get_push_config(
1012-
&self,
1013-
_params: &ServiceParams,
1014-
_req: GetTaskPushNotificationConfigRequest,
1015-
) -> Result<TaskPushNotificationConfig, A2AError> {
1016-
unreachable!()
1017-
}
1018-
1019-
async fn list_push_configs(
1020-
&self,
1021-
_params: &ServiceParams,
1022-
_req: ListTaskPushNotificationConfigsRequest,
1023-
) -> Result<ListTaskPushNotificationConfigsResponse, A2AError> {
1024-
unreachable!()
1025-
}
1026-
1027-
async fn delete_push_config(
1028-
&self,
1029-
_params: &ServiceParams,
1030-
_req: DeleteTaskPushNotificationConfigRequest,
1031-
) -> Result<(), A2AError> {
1032-
unreachable!()
1033-
}
1034-
1035-
async fn get_extended_agent_card(
1036-
&self,
1037-
_params: &ServiceParams,
1038-
_req: GetExtendedAgentCardRequest,
1039-
) -> Result<AgentCard, A2AError> {
1040-
unreachable!()
1041-
}
1042-
}
1043-
1044886
#[test]
1045887
fn websocket_router_constructs_with_request_handler() {
1046888
let _router = websocket_router(make_handler());
@@ -1230,7 +1072,7 @@ mod tests {
12301072

12311073
#[tokio::test]
12321074
async fn handle_text_frame_invalid_json_sends_error_and_close() {
1233-
let handler = Arc::new(StubHandler);
1075+
let handler = Arc::new(StubHandler::default());
12341076
let params = Arc::new(ServiceParams::new());
12351077
let streams = Arc::new(Mutex::new(HashMap::new()));
12361078
let (out_tx, mut out_rx) = mpsc::channel(OUTBOUND_BUFFER_CAPACITY);
@@ -1252,7 +1094,7 @@ mod tests {
12521094

12531095
#[tokio::test]
12541096
async fn handle_text_frame_empty_id_sends_invalid_request() {
1255-
let handler = Arc::new(StubHandler);
1097+
let handler = Arc::new(StubHandler::default());
12561098
let params = Arc::new(ServiceParams::new());
12571099
let streams = Arc::new(Mutex::new(HashMap::new()));
12581100
let (out_tx, mut out_rx) = mpsc::channel(OUTBOUND_BUFFER_CAPACITY);
@@ -1275,7 +1117,7 @@ mod tests {
12751117

12761118
#[tokio::test]
12771119
async fn handle_text_frame_missing_method_sends_invalid_request() {
1278-
let handler = Arc::new(StubHandler);
1120+
let handler = Arc::new(StubHandler::default());
12791121
let params = Arc::new(ServiceParams::new());
12801122
let streams = Arc::new(Mutex::new(HashMap::new()));
12811123
let (out_tx, mut out_rx) = mpsc::channel(OUTBOUND_BUFFER_CAPACITY);
@@ -1297,7 +1139,7 @@ mod tests {
12971139

12981140
#[tokio::test]
12991141
async fn handle_text_frame_unknown_method_sends_method_not_found() {
1300-
let handler = Arc::new(StubHandler);
1142+
let handler = Arc::new(StubHandler::default());
13011143
let params = Arc::new(ServiceParams::new());
13021144
let streams = Arc::new(Mutex::new(HashMap::new()));
13031145
let (out_tx, mut out_rx) = mpsc::channel(OUTBOUND_BUFFER_CAPACITY);
@@ -1320,7 +1162,7 @@ mod tests {
13201162

13211163
#[tokio::test]
13221164
async fn handle_text_frame_cancel_stream_removes_registered_stream() {
1323-
let handler = Arc::new(StubHandler);
1165+
let handler = Arc::new(StubHandler::default());
13241166
let params = Arc::new(ServiceParams::new());
13251167
let streams = Arc::new(Mutex::new(HashMap::new()));
13261168
let (cancel_tx, cancel_rx) = oneshot::channel();
@@ -1342,7 +1184,7 @@ mod tests {
13421184

13431185
#[tokio::test]
13441186
async fn handle_text_frame_cancel_unknown_stream_does_not_emit_stream_end() {
1345-
let handler = Arc::new(StubHandler);
1187+
let handler = Arc::new(StubHandler::default());
13461188
let params = Arc::new(ServiceParams::new());
13471189
let streams = Arc::new(Mutex::new(HashMap::new()));
13481190
let (out_tx, mut out_rx) = mpsc::channel(OUTBOUND_BUFFER_CAPACITY);
@@ -1361,7 +1203,7 @@ mod tests {
13611203

13621204
#[tokio::test]
13631205
async fn dispatch_unary_covers_all_supported_methods() {
1364-
let handler = Arc::new(StubHandler);
1206+
let handler = Arc::new(StubHandler::default());
13651207
let params = ServiceParams::new();
13661208
let msg_req = SendMessageRequest {
13671209
message: sample_message(),
@@ -1510,7 +1352,7 @@ mod tests {
15101352

15111353
#[tokio::test]
15121354
async fn dispatch_unary_unknown_method_returns_method_not_found() {
1513-
let handler = Arc::new(StubHandler);
1355+
let handler = Arc::new(StubHandler::default());
15141356
let err = dispatch_unary("Nope", &handler, &ServiceParams::new(), Value::Null)
15151357
.await
15161358
.unwrap_err();
@@ -1519,7 +1361,7 @@ mod tests {
15191361

15201362
#[tokio::test]
15211363
async fn run_unary_request_emits_error_for_bad_params() {
1522-
let handler = Arc::new(StubHandler);
1364+
let handler = Arc::new(StubHandler::default());
15231365
let (out_tx, mut out_rx) = mpsc::channel(OUTBOUND_BUFFER_CAPACITY);
15241366

15251367
run_unary_request(
@@ -1542,7 +1384,7 @@ mod tests {
15421384

15431385
#[tokio::test]
15441386
async fn run_unary_request_emits_close_for_fatal_error() {
1545-
let handler = Arc::new(FatalHandler);
1387+
let handler = Arc::new(StubHandler::fatal_send_message());
15461388
let (out_tx, mut out_rx) = mpsc::channel(OUTBOUND_BUFFER_CAPACITY);
15471389
let req = SendMessageRequest {
15481390
message: sample_message(),
@@ -1576,7 +1418,7 @@ mod tests {
15761418

15771419
#[tokio::test]
15781420
async fn run_streaming_request_emits_event_and_stream_end() {
1579-
let handler = Arc::new(StubHandler);
1421+
let handler = Arc::new(StubHandler::default());
15801422
let streams = Arc::new(Mutex::new(HashMap::new()));
15811423
let (out_tx, mut out_rx) = mpsc::channel(OUTBOUND_BUFFER_CAPACITY);
15821424
let req = SendMessageRequest {
@@ -1607,7 +1449,7 @@ mod tests {
16071449

16081450
#[tokio::test]
16091451
async fn run_streaming_request_emits_stream_end_after_cancellation() {
1610-
let handler = Arc::new(PendingStreamHandler);
1452+
let handler = Arc::new(StubHandler::pending_stream());
16111453
let streams = Arc::new(Mutex::new(HashMap::new()));
16121454
let (out_tx, mut out_rx) = mpsc::channel(OUTBOUND_BUFFER_CAPACITY);
16131455
let req = SendMessageRequest {
@@ -1649,7 +1491,7 @@ mod tests {
16491491

16501492
#[tokio::test]
16511493
async fn run_streaming_request_emits_error_for_stream_item_error() {
1652-
let handler = Arc::new(StubHandler);
1494+
let handler = Arc::new(StubHandler::default());
16531495
let streams = Arc::new(Mutex::new(HashMap::new()));
16541496
let (out_tx, mut out_rx) = mpsc::channel(OUTBOUND_BUFFER_CAPACITY);
16551497

@@ -1677,7 +1519,7 @@ mod tests {
16771519

16781520
#[tokio::test]
16791521
async fn run_streaming_request_emits_error_for_bad_stream_params() {
1680-
let handler = Arc::new(StubHandler);
1522+
let handler = Arc::new(StubHandler::default());
16811523
let streams = Arc::new(Mutex::new(HashMap::new()));
16821524
let (out_tx, mut out_rx) = mpsc::channel(OUTBOUND_BUFFER_CAPACITY);
16831525

0 commit comments

Comments
 (0)