improve test readability

codex-rs/mcp-server/tests/stream_conversation.rs

@@ -38,15 +38,19 @@ async fn test_connect_then_send_receives_initial_state_and_notifications() {
         .expect("create conversation");
 
     // Connect the stream
-    let (_stream_req, params) = mcp
+    let (stream_req, params) = mcp
         .connect_stream_and_expect_initial_state(&conv_id)
         .await
         .expect("initial_state params");
-    assert_eq!(
-        params["_meta"]["conversationId"].as_str(),
-        Some(conv_id.as_str())
-    );
-    assert_eq!(params["initial_state"], json!({ "events": [] }));
+    let expected_params = json!({
+        "_meta": {
+            "conversationId": conv_id.as_str(),
+        },
+        "initial_state": {
+            "events": []
+        }
+    });
+    assert_eq!(params, expected_params);
 
     // Send a message and expect a subsequent notification (non-initial_state)
     mcp.send_user_message_and_wait_ok("Hello there", &conv_id)
@@ -55,10 +59,13 @@ async fn test_connect_then_send_receives_initial_state_and_notifications() {
 
     // Read until we see an event notification (new schema example: agent_message)
     let params = mcp.wait_for_agent_message().await.expect("agent message");
-    assert_eq!(
-        params["msg"],
-        json!({ "type": "agent_message", "message": "Done" })
-    );
+    let expected_params = json!({
+        "msg": {
+            "type": "agent_message",
+            "message": "Done"
+        }
+    });
+    assert_eq!(params, expected_params);
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
@@ -98,28 +105,34 @@ async fn test_send_then_connect_receives_initial_state_with_message() {
     let events = params["initial_state"]["events"]
         .as_array()
         .expect("events array");
-    let mut agent_events: Vec<_> = events
-        .iter()
-        .filter(|ev| ev["msg"]["type"].as_str() == Some("agent_message"))
-        .cloned()
-        .collect();
-    if agent_events.is_empty() {
-        // Fallback to live notification if not present in initial state, then assert the full event list
+    if !events.iter().any(|ev| {
+        ev.get("msg")
+            .and_then(|m| m.get("type"))
+            .and_then(|t| t.as_str())
+            == Some("agent_message")
+            && ev
+                .get("msg")
+                .and_then(|m| m.get("message"))
+                .and_then(|t| t.as_str())
+                == Some("Done")
+    }) {
+        // Fallback to live notification if not present in initial state
         let note: JSONRPCNotification = timeout(
             DEFAULT_READ_TIMEOUT,
             mcp.read_stream_until_notification_method("agent_message"),
         )
         .await
-        .expect("agent_message note timeout")
-        .expect("agent_message note err");
-        let p = note.params.expect("params");
-        agent_events.push(json!({ "msg": { "type": "agent_message", "message": p["msg"]["message"].as_str().expect("message str") } }));
+        .expect("event note timeout")
+        .expect("event note err");
+        let params = note.params.expect("params");
+        let expected_params = json!({
+            "msg": {
+                "type": "agent_message",
+                "message": "Done"
+            }
+        });
+        assert_eq!(params, expected_params);
     }
-    let expected = vec![json!({ "msg": { "type": "agent_message", "message": "Done" } })];
-    assert_eq!(
-        agent_events, expected,
-        "initial_state agent_message events should match exactly"
-    );
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
@@ -189,19 +202,43 @@ async fn test_cancel_stream_then_reconnect_catches_up_initial_state() {
     let events = params["initial_state"]["events"]
         .as_array()
         .expect("events array");
-    let agent_events: Vec<_> = events
-        .iter()
-        .filter(|ev| ev["msg"]["type"].as_str() == Some("agent_message"))
-        .cloned()
-        .collect();
     let expected = vec![
-        json!({ "msg": { "type": "agent_message", "message": "Done 1" } }),
-        json!({ "msg": { "type": "agent_message", "message": "Done 2" } }),
+        json!({
+            "msg": {
+                "type": "task_started",
+            },
+        }),
+        json!({
+            "msg": {
+                "message": "Done 1",
+                "type": "agent_message",
+            },
+        }),
+        json!({
+            "msg": {
+                "last_agent_message": "Done 1",
+                "type": "task_complete",
+            },
+        }),
+        json!({
+            "msg": {
+                "type": "task_started",
+            },
+        }),
+        json!({
+            "msg": {
+                "message": "Done 2",
+                "type": "agent_message",
+            },
+        }),
+        json!({
+            "msg": {
+                "last_agent_message": "Done 2",
+                "type": "task_complete",
+            },
+        }),
     ];
-    assert_eq!(
-        agent_events, expected,
-        "initial_state agent_message events should match exactly"
-    );
+    assert_eq!(*events, expected);
     drop(server);
 }