Skip to content

Commit cffc5dc

Browse files
yordisyordis
committed
chore(acp-nats-agent): strengthen dispatch test assertions
Signed-off-by: Yordis Prieto <yordis.prieto@gmail.com>
1 parent 490eb99 commit cffc5dc

1 file changed

Lines changed: 94 additions & 76 deletions

File tree

‎rsworkspace/crates/acp-nats-agent/src/connection.rs‎

Lines changed: 94 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -695,7 +695,7 @@ mod tests {
695695
use super::*;
696696
use agent_client_protocol::{
697697
AuthenticateResponse, Error as AcpError, ErrorCode, InitializeResponse, LogoutResponse,
698-
PromptResponse, StopReason,
698+
NewSessionResponse, PromptResponse, StopReason,
699699
};
700700
use std::cell::RefCell;
701701
use trogon_nats::MockNatsClient;
@@ -866,7 +866,9 @@ mod tests {
866866

867867
#[tokio::test]
868868
async fn dispatch_logout_publishes_response() {
869-
assert_dispatch_publishes("acp.agent.logout", &LogoutRequest::new()).await;
869+
let (nats, _) = dispatch("acp.agent.logout", &LogoutRequest::new(), Some("_INBOX.r")).await;
870+
assert_eq!(nats.published_messages(), vec!["_INBOX.r"]);
871+
let _: LogoutResponse = published_response(&nats);
870872
}
871873

872874
#[tokio::test]
@@ -878,7 +880,7 @@ mod tests {
878880
)
879881
.await;
880882

881-
assert_eq!(agent.cancelled.borrow().len(), 1);
883+
assert_eq!(agent.cancelled.borrow().as_slice(), ["s1"]);
882884
assert!(nats.published_messages().is_empty());
883885
}
884886

@@ -926,6 +928,11 @@ mod tests {
926928
)
927929
.await;
928930
assert_eq!(nats.published_messages(), vec!["_INBOX.specific"]);
931+
let response: InitializeResponse = published_response(&nats);
932+
assert_eq!(
933+
response.protocol_version,
934+
agent_client_protocol::ProtocolVersion::V0
935+
);
929936
}
930937

931938
#[test]
@@ -984,6 +991,8 @@ mod tests {
984991
)
985992
.await;
986993
assert_eq!(nats.published_messages(), vec!["_INBOX.ext"]);
994+
let value: serde_json::Value = published_response(&nats);
995+
assert!(value.is_null());
987996
}
988997

989998
#[tokio::test]
@@ -997,19 +1006,30 @@ mod tests {
9971006
assert!(nats.published_messages().is_empty());
9981007
}
9991008

1000-
async fn assert_dispatch_publishes<T: serde::Serialize>(subject: &str, args: &T) {
1009+
async fn assert_dispatch_method_not_found<T: serde::Serialize>(subject: &str, args: &T) {
10011010
let (nats, _) = dispatch(subject, args, Some("_INBOX.r")).await;
10021011
assert_eq!(nats.published_messages(), vec!["_INBOX.r"]);
1012+
let error: AcpError = published_response(&nats);
1013+
assert_eq!(error.code, ErrorCode::MethodNotFound);
10031014
}
10041015

10051016
#[tokio::test]
10061017
async fn dispatch_new_session_publishes_response() {
1007-
assert_dispatch_publishes("acp.agent.session.new", &NewSessionRequest::new("/tmp")).await;
1018+
let (nats, _) = dispatch(
1019+
"acp.agent.session.new",
1020+
&NewSessionRequest::new("/tmp"),
1021+
Some("_INBOX.r"),
1022+
)
1023+
.await;
1024+
1025+
assert_eq!(nats.published_messages(), vec!["_INBOX.r"]);
1026+
let response: NewSessionResponse = published_response(&nats);
1027+
assert_eq!(response.session_id.to_string(), "sess-1");
10081028
}
10091029

10101030
#[tokio::test]
10111031
async fn dispatch_session_load_publishes_response() {
1012-
assert_dispatch_publishes(
1032+
assert_dispatch_method_not_found(
10131033
"acp.session.s1.agent.load",
10141034
&LoadSessionRequest::new("s1", "/tmp"),
10151035
)
@@ -1018,12 +1038,13 @@ mod tests {
10181038

10191039
#[tokio::test]
10201040
async fn dispatch_list_sessions_publishes_response() {
1021-
assert_dispatch_publishes("acp.agent.session.list", &ListSessionsRequest::new()).await;
1041+
assert_dispatch_method_not_found("acp.agent.session.list", &ListSessionsRequest::new())
1042+
.await;
10221043
}
10231044

10241045
#[tokio::test]
10251046
async fn dispatch_set_session_mode_publishes_response() {
1026-
assert_dispatch_publishes(
1047+
assert_dispatch_method_not_found(
10271048
"acp.session.s1.agent.set_mode",
10281049
&SetSessionModeRequest::new("s1", "code"),
10291050
)
@@ -1032,7 +1053,7 @@ mod tests {
10321053

10331054
#[tokio::test]
10341055
async fn dispatch_set_session_config_option_publishes_response() {
1035-
assert_dispatch_publishes(
1056+
assert_dispatch_method_not_found(
10361057
"acp.session.s1.agent.set_config_option",
10371058
&SetSessionConfigOptionRequest::new("s1", "key", "val"),
10381059
)
@@ -1041,7 +1062,7 @@ mod tests {
10411062

10421063
#[tokio::test]
10431064
async fn dispatch_set_session_model_publishes_response() {
1044-
assert_dispatch_publishes(
1065+
assert_dispatch_method_not_found(
10451066
"acp.session.s1.agent.set_model",
10461067
&SetSessionModelRequest::new("s1", "gpt-4"),
10471068
)
@@ -1050,7 +1071,7 @@ mod tests {
10501071

10511072
#[tokio::test]
10521073
async fn dispatch_fork_session_publishes_response() {
1053-
assert_dispatch_publishes(
1074+
assert_dispatch_method_not_found(
10541075
"acp.session.s1.agent.fork",
10551076
&ForkSessionRequest::new("s1", "/tmp"),
10561077
)
@@ -1059,7 +1080,7 @@ mod tests {
10591080

10601081
#[tokio::test]
10611082
async fn dispatch_resume_session_publishes_response() {
1062-
assert_dispatch_publishes(
1083+
assert_dispatch_method_not_found(
10631084
"acp.session.s1.agent.resume",
10641085
&ResumeSessionRequest::new("s1", "/tmp"),
10651086
)
@@ -1068,7 +1089,7 @@ mod tests {
10681089

10691090
#[tokio::test]
10701091
async fn dispatch_close_session_publishes_response() {
1071-
assert_dispatch_publishes(
1092+
assert_dispatch_method_not_found(
10721093
"acp.session.s1.agent.close",
10731094
&CloseSessionRequest::new("s1"),
10741095
)
@@ -1251,7 +1272,12 @@ mod tests {
12511272
tokio::task::yield_now().await;
12521273
tokio::task::yield_now().await;
12531274

1254-
assert!(!nats.published_messages().is_empty());
1275+
assert_eq!(nats.published_messages(), vec!["_INBOX.serve"]);
1276+
let response: InitializeResponse = published_response(&nats);
1277+
assert_eq!(
1278+
response.protocol_version,
1279+
agent_client_protocol::ProtocolVersion::V0
1280+
);
12551281
})
12561282
.await;
12571283
}
@@ -1301,7 +1327,7 @@ mod tests {
13011327
tokio::task::yield_now().await;
13021328
tokio::task::yield_now().await;
13031329

1304-
assert!(!nats.published_messages().is_empty());
1330+
assert_js_response_method_not_found(&nats, "acp.session.s1.agent.response.req-1");
13051331
})
13061332
.await;
13071333
}
@@ -1364,18 +1390,6 @@ mod tests {
13641390
.await;
13651391
}
13661392

1367-
#[tokio::test]
1368-
async fn dispatch_js_message_success_acks() {
1369-
let nats = MockNatsClient::new();
1370-
let agent = MockAgent::new();
1371-
let payload = serialize(&LoadSessionRequest::new("s1", "/tmp"));
1372-
let js_msg = make_js_msg("acp.session.s1.agent.load", &payload, None);
1373-
1374-
dispatch_js_message(js_msg, &agent, &nats, &test_prefix()).await;
1375-
1376-
assert!(!nats.published_messages().is_empty());
1377-
}
1378-
13791393
#[tokio::test]
13801394
async fn dispatch_js_message_unknown_subject_terms() {
13811395
let nats = MockNatsClient::new();
@@ -1395,9 +1409,11 @@ mod tests {
13951409

13961410
dispatch_js_message(js_msg, &agent, &nats, &test_prefix()).await;
13971411

1398-
let payloads = nats.published_payloads();
1399-
assert_eq!(payloads.len(), 1);
1400-
let error: agent_client_protocol::Error = serde_json::from_slice(&payloads[0]).unwrap();
1412+
assert_eq!(
1413+
nats.published_messages(),
1414+
vec!["acp.session.s1.agent.response.req-1"]
1415+
);
1416+
let error: AcpError = published_response(&nats);
14011417
assert_eq!(error.code, ErrorCode::InvalidParams);
14021418
}
14031419

@@ -1528,14 +1544,12 @@ mod tests {
15281544

15291545
dispatch_js_message(js_msg, &agent, &nats, &test_prefix()).await;
15301546

1531-
let subjects = nats.published_messages();
1532-
assert!(
1533-
subjects
1534-
.iter()
1535-
.any(|s| s.starts_with("acp.session.s1.agent.prompt.response.")),
1536-
"expected prompt.response subject, got: {:?}",
1537-
subjects
1547+
assert_eq!(
1548+
nats.published_messages(),
1549+
vec!["acp.session.s1.agent.prompt.response.req-1"]
15381550
);
1551+
let response: PromptResponse = published_response(&nats);
1552+
assert_eq!(response.stop_reason, StopReason::EndTurn);
15391553
}
15401554

15411555
#[tokio::test]
@@ -1547,14 +1561,7 @@ mod tests {
15471561

15481562
dispatch_js_message(js_msg, &agent, &nats, &test_prefix()).await;
15491563

1550-
let subjects = nats.published_messages();
1551-
assert!(
1552-
subjects
1553-
.iter()
1554-
.any(|s| s.starts_with("acp.session.s1.agent.response.")),
1555-
"expected response subject, got: {:?}",
1556-
subjects
1557-
);
1564+
assert_js_response_method_not_found(&nats, "acp.session.s1.agent.response.req-1");
15581565
}
15591566

15601567
#[tokio::test]
@@ -1610,8 +1617,12 @@ mod tests {
16101617
tokio::task::yield_now().await;
16111618
tokio::task::yield_now().await;
16121619

1613-
assert_eq!(nats.published_messages().len(), 1);
1614-
assert_eq!(nats.published_messages()[0], "_INBOX.serve");
1620+
assert_eq!(nats.published_messages(), vec!["_INBOX.serve"]);
1621+
let response: InitializeResponse = published_response(&nats);
1622+
assert_eq!(
1623+
response.protocol_version,
1624+
agent_client_protocol::ProtocolVersion::V0
1625+
);
16151626
})
16161627
.await;
16171628
}
@@ -1688,87 +1699,85 @@ mod tests {
16881699

16891700
dispatch_js_message(js_msg, &agent, &nats, &test_prefix()).await;
16901701

1691-
assert_eq!(agent.cancelled.borrow().len(), 1);
1702+
assert_eq!(agent.cancelled.borrow().as_slice(), ["s1"]);
1703+
}
1704+
1705+
fn assert_js_response_method_not_found(nats: &MockNatsClient, expected_subject: &str) {
1706+
assert_eq!(nats.published_messages(), vec![expected_subject]);
1707+
let error: AcpError = published_response(nats);
1708+
assert_eq!(error.code, ErrorCode::MethodNotFound);
16921709
}
16931710

16941711
#[tokio::test]
16951712
async fn dispatch_js_message_set_mode() {
16961713
let nats = MockNatsClient::new();
16971714
let agent = MockAgent::new();
16981715
let payload = serialize(&SetSessionModeRequest::new("s1", "code"));
1699-
let js_msg = make_js_msg("acp.session.s1.agent.set_mode", &payload, Some("_INBOX.r"));
1716+
let js_msg = make_js_msg("acp.session.s1.agent.set_mode", &payload, None);
17001717

17011718
dispatch_js_message(js_msg, &agent, &nats, &test_prefix()).await;
17021719

1703-
assert!(!nats.published_messages().is_empty());
1720+
assert_js_response_method_not_found(&nats, "acp.session.s1.agent.response.req-1");
17041721
}
17051722

17061723
#[tokio::test]
17071724
async fn dispatch_js_message_close_session() {
17081725
let nats = MockNatsClient::new();
17091726
let agent = MockAgent::new();
17101727
let payload = serialize(&CloseSessionRequest::new("s1"));
1711-
let js_msg = make_js_msg("acp.session.s1.agent.close", &payload, Some("_INBOX.r"));
1728+
let js_msg = make_js_msg("acp.session.s1.agent.close", &payload, None);
17121729

17131730
dispatch_js_message(js_msg, &agent, &nats, &test_prefix()).await;
17141731

1715-
assert!(!nats.published_messages().is_empty());
1732+
assert_js_response_method_not_found(&nats, "acp.session.s1.agent.response.req-1");
17161733
}
17171734

17181735
#[tokio::test]
17191736
async fn dispatch_js_message_fork_session() {
17201737
let nats = MockNatsClient::new();
17211738
let agent = MockAgent::new();
17221739
let payload = serialize(&ForkSessionRequest::new("s1", "/tmp"));
1723-
let js_msg = make_js_msg("acp.session.s1.agent.fork", &payload, Some("_INBOX.r"));
1740+
let js_msg = make_js_msg("acp.session.s1.agent.fork", &payload, None);
17241741

17251742
dispatch_js_message(js_msg, &agent, &nats, &test_prefix()).await;
17261743

1727-
assert!(!nats.published_messages().is_empty());
1744+
assert_js_response_method_not_found(&nats, "acp.session.s1.agent.response.req-1");
17281745
}
17291746

17301747
#[tokio::test]
17311748
async fn dispatch_js_message_set_config_option() {
17321749
let nats = MockNatsClient::new();
17331750
let agent = MockAgent::new();
17341751
let payload = serialize(&SetSessionConfigOptionRequest::new("s1", "key", "val"));
1735-
let js_msg = make_js_msg(
1736-
"acp.session.s1.agent.set_config_option",
1737-
&payload,
1738-
Some("_INBOX.r"),
1739-
);
1752+
let js_msg = make_js_msg("acp.session.s1.agent.set_config_option", &payload, None);
1753+
17401754
dispatch_js_message(js_msg, &agent, &nats, &test_prefix()).await;
1741-
assert!(!nats.published_messages().is_empty());
1755+
1756+
assert_js_response_method_not_found(&nats, "acp.session.s1.agent.response.req-1");
17421757
}
17431758

17441759
#[tokio::test]
17451760
async fn dispatch_js_message_set_model() {
17461761
let nats = MockNatsClient::new();
17471762
let agent = MockAgent::new();
17481763
let payload = serialize(&SetSessionModelRequest::new("s1", "gpt-4"));
1749-
let js_msg = make_js_msg("acp.session.s1.agent.set_model", &payload, Some("_INBOX.r"));
1764+
let js_msg = make_js_msg("acp.session.s1.agent.set_model", &payload, None);
1765+
17501766
dispatch_js_message(js_msg, &agent, &nats, &test_prefix()).await;
1751-
assert!(!nats.published_messages().is_empty());
1767+
1768+
assert_js_response_method_not_found(&nats, "acp.session.s1.agent.response.req-1");
17521769
}
17531770

17541771
#[tokio::test]
17551772
async fn dispatch_js_message_resume_session() {
17561773
let nats = MockNatsClient::new();
17571774
let agent = MockAgent::new();
17581775
let payload = serialize(&ResumeSessionRequest::new("s1", "/tmp"));
1759-
let js_msg = make_js_msg("acp.session.s1.agent.resume", &payload, Some("_INBOX.r"));
1760-
dispatch_js_message(js_msg, &agent, &nats, &test_prefix()).await;
1761-
assert!(!nats.published_messages().is_empty());
1762-
}
1776+
let js_msg = make_js_msg("acp.session.s1.agent.resume", &payload, None);
17631777

1764-
#[tokio::test]
1765-
async fn dispatch_js_message_prompt() {
1766-
let nats = MockNatsClient::new();
1767-
let agent = MockAgent::new();
1768-
let payload = serialize(&PromptRequest::new("s1", vec![]));
1769-
let js_msg = make_js_msg("acp.session.s1.agent.prompt", &payload, Some("_INBOX.r"));
17701778
dispatch_js_message(js_msg, &agent, &nats, &test_prefix()).await;
1771-
assert!(!nats.published_messages().is_empty());
1779+
1780+
assert_js_response_method_not_found(&nats, "acp.session.s1.agent.response.req-1");
17721781
}
17731782

17741783
#[tokio::test]
@@ -1900,7 +1909,12 @@ mod tests {
19001909
})
19011910
.await;
19021911
assert!(result.is_ok());
1903-
assert!(!nats.published_messages().is_empty());
1912+
assert_eq!(nats.published_messages(), vec!["_INBOX.1"]);
1913+
let response: InitializeResponse = published_response(&nats);
1914+
assert_eq!(
1915+
response.protocol_version,
1916+
agent_client_protocol::ProtocolVersion::V0
1917+
);
19041918
}
19051919

19061920
#[tokio::test]
@@ -1934,7 +1948,9 @@ mod tests {
19341948
let js_msg = make_js_msg("acp.agent.initialize", &payload, Some("_INBOX.1"));
19351949
let result = handle_request_with_keepalive(&msg, &nats, &js_msg, init_handler_error).await;
19361950
assert!(result.is_ok());
1937-
assert!(!nats.published_messages().is_empty());
1951+
assert_eq!(nats.published_messages(), vec!["_INBOX.1"]);
1952+
let error: AcpError = published_response(&nats);
1953+
assert_eq!(error.code, ErrorCode::InternalError);
19381954
}
19391955

19401956
#[tokio::test(start_paused = true)]
@@ -1984,6 +2000,8 @@ mod tests {
19842000
})
19852001
.await;
19862002
assert!(result.is_ok());
1987-
assert!(!nats.published_messages().is_empty());
2003+
assert_eq!(nats.published_messages(), vec!["_INBOX.1"]);
2004+
let error: AcpError = published_response(&nats);
2005+
assert_eq!(error.code, ErrorCode::MethodNotFound);
19882006
}
19892007
}

0 commit comments

Comments
 (0)