Skip to content

Commit 51905cd

Browse files
hayemaxiKonippi
hayemaxi
authored and
Konippi
committed
test(tools): add integration tests for /tools + chat usage (aws#1081)
aws#1080 - Test /tools usage on how tools are executed - Test running multiple tools at in one Q response
1 parent 6e01028 commit 51905cd

File tree

1 file changed

+294
-0
lines changed
  • crates/q_cli/src/cli/chat

1 file changed

+294
-0
lines changed

‎crates/q_cli/src/cli/chat/mod.rs

+294
Original file line numberDiff line numberDiff line change
@@ -2076,6 +2076,300 @@ mod tests {
20762076
assert_eq!(ctx.fs().read_to_string("/file.txt").await.unwrap(), "Hello, world!\n");
20772077
}
20782078

2079+
#[tokio::test]
2080+
async fn test_flow_tool_permissions() {
2081+
let _ = tracing_subscriber::fmt::try_init();
2082+
let ctx = Context::builder().with_test_home().await.unwrap().build_fake();
2083+
let test_client = create_stream(serde_json::json!([
2084+
[
2085+
"Ok",
2086+
{
2087+
"tool_use_id": "1",
2088+
"name": "fs_write",
2089+
"args": {
2090+
"command": "create",
2091+
"file_text": "Hello, world!",
2092+
"path": "/file1.txt",
2093+
}
2094+
}
2095+
],
2096+
[
2097+
"Done",
2098+
],
2099+
[
2100+
"Ok",
2101+
{
2102+
"tool_use_id": "1",
2103+
"name": "fs_write",
2104+
"args": {
2105+
"command": "create",
2106+
"file_text": "Hello, world!",
2107+
"path": "/file2.txt",
2108+
}
2109+
}
2110+
],
2111+
[
2112+
"Done",
2113+
],
2114+
[
2115+
"Ok",
2116+
{
2117+
"tool_use_id": "1",
2118+
"name": "fs_write",
2119+
"args": {
2120+
"command": "create",
2121+
"file_text": "Hello, world!",
2122+
"path": "/file3.txt",
2123+
}
2124+
}
2125+
],
2126+
[
2127+
"Done",
2128+
],
2129+
[
2130+
"Ok",
2131+
{
2132+
"tool_use_id": "1",
2133+
"name": "fs_write",
2134+
"args": {
2135+
"command": "create",
2136+
"file_text": "Hello, world!",
2137+
"path": "/file4.txt",
2138+
}
2139+
}
2140+
],
2141+
[
2142+
"Ok, I won't make it.",
2143+
],
2144+
[
2145+
"Ok",
2146+
{
2147+
"tool_use_id": "1",
2148+
"name": "fs_write",
2149+
"args": {
2150+
"command": "create",
2151+
"file_text": "Hello, world!",
2152+
"path": "/file5.txt",
2153+
}
2154+
}
2155+
],
2156+
[
2157+
"Done",
2158+
],
2159+
[
2160+
"Ok",
2161+
{
2162+
"tool_use_id": "1",
2163+
"name": "fs_write",
2164+
"args": {
2165+
"command": "create",
2166+
"file_text": "Hello, world!",
2167+
"path": "/file6.txt",
2168+
}
2169+
}
2170+
],
2171+
[
2172+
"Ok, I won't make it.",
2173+
],
2174+
]));
2175+
2176+
ChatContext::new(
2177+
Arc::clone(&ctx),
2178+
Settings::new_fake(),
2179+
std::io::stdout(),
2180+
None,
2181+
InputSource::new_mock(vec![
2182+
"/tools".to_string(),
2183+
"/tools help".to_string(),
2184+
"create a new file".to_string(),
2185+
"y".to_string(),
2186+
"create a new file".to_string(),
2187+
"t".to_string(),
2188+
"create a new file".to_string(), // should make without prompting due to 't'
2189+
"/tools untrust fs_write".to_string(),
2190+
"create a file".to_string(), // prompt again due to untrust
2191+
"n".to_string(), // cancel
2192+
"/tools trust fs_write".to_string(),
2193+
"create a file".to_string(), // again without prompting due to '/tools trust'
2194+
"/tools reset".to_string(),
2195+
"create a file".to_string(), // prompt again due to reset
2196+
"n".to_string(), // cancel
2197+
"exit".to_string(),
2198+
]),
2199+
true,
2200+
test_client,
2201+
|| Some(80),
2202+
None,
2203+
load_tools().expect("Tools failed to load."),
2204+
ToolPermissions::new(0),
2205+
)
2206+
.await
2207+
.unwrap()
2208+
.try_chat()
2209+
.await
2210+
.unwrap();
2211+
2212+
assert_eq!(ctx.fs().read_to_string("/file2.txt").await.unwrap(), "Hello, world!\n");
2213+
assert_eq!(ctx.fs().read_to_string("/file3.txt").await.unwrap(), "Hello, world!\n");
2214+
assert!(!ctx.fs().exists("/file4.txt"));
2215+
assert_eq!(ctx.fs().read_to_string("/file5.txt").await.unwrap(), "Hello, world!\n");
2216+
assert!(!ctx.fs().exists("/file6.txt"));
2217+
}
2218+
2219+
#[tokio::test]
2220+
async fn test_flow_multiple_tools() {
2221+
let _ = tracing_subscriber::fmt::try_init();
2222+
let ctx = Context::builder().with_test_home().await.unwrap().build_fake();
2223+
let test_client = create_stream(serde_json::json!([
2224+
[
2225+
"Sure, I'll create a file for you",
2226+
{
2227+
"tool_use_id": "1",
2228+
"name": "fs_write",
2229+
"args": {
2230+
"command": "create",
2231+
"file_text": "Hello, world!",
2232+
"path": "/file1.txt",
2233+
}
2234+
},
2235+
{
2236+
"tool_use_id": "2",
2237+
"name": "fs_write",
2238+
"args": {
2239+
"command": "create",
2240+
"file_text": "Hello, world!",
2241+
"path": "/file2.txt",
2242+
}
2243+
}
2244+
],
2245+
[
2246+
"Done",
2247+
],
2248+
[
2249+
"Sure, I'll create a file for you",
2250+
{
2251+
"tool_use_id": "1",
2252+
"name": "fs_write",
2253+
"args": {
2254+
"command": "create",
2255+
"file_text": "Hello, world!",
2256+
"path": "/file3.txt",
2257+
}
2258+
},
2259+
{
2260+
"tool_use_id": "2",
2261+
"name": "fs_write",
2262+
"args": {
2263+
"command": "create",
2264+
"file_text": "Hello, world!",
2265+
"path": "/file4.txt",
2266+
}
2267+
}
2268+
],
2269+
[
2270+
"Done",
2271+
],
2272+
]));
2273+
2274+
ChatContext::new(
2275+
Arc::clone(&ctx),
2276+
Settings::new_fake(),
2277+
std::io::stdout(),
2278+
None,
2279+
InputSource::new_mock(vec![
2280+
"create 2 new files parallel".to_string(),
2281+
"t".to_string(),
2282+
"/tools reset".to_string(),
2283+
"create 2 new files parallel".to_string(),
2284+
"y".to_string(),
2285+
"y".to_string(),
2286+
"exit".to_string(),
2287+
]),
2288+
true,
2289+
test_client,
2290+
|| Some(80),
2291+
None,
2292+
load_tools().expect("Tools failed to load."),
2293+
ToolPermissions::new(0),
2294+
)
2295+
.await
2296+
.unwrap()
2297+
.try_chat()
2298+
.await
2299+
.unwrap();
2300+
2301+
assert_eq!(ctx.fs().read_to_string("/file1.txt").await.unwrap(), "Hello, world!\n");
2302+
assert_eq!(ctx.fs().read_to_string("/file2.txt").await.unwrap(), "Hello, world!\n");
2303+
assert_eq!(ctx.fs().read_to_string("/file3.txt").await.unwrap(), "Hello, world!\n");
2304+
assert_eq!(ctx.fs().read_to_string("/file4.txt").await.unwrap(), "Hello, world!\n");
2305+
}
2306+
2307+
#[tokio::test]
2308+
async fn test_flow_tools_trust_all() {
2309+
let _ = tracing_subscriber::fmt::try_init();
2310+
let ctx = Context::builder().with_test_home().await.unwrap().build_fake();
2311+
let test_client = create_stream(serde_json::json!([
2312+
[
2313+
"Sure, I'll create a file for you",
2314+
{
2315+
"tool_use_id": "1",
2316+
"name": "fs_write",
2317+
"args": {
2318+
"command": "create",
2319+
"file_text": "Hello, world!",
2320+
"path": "/file1.txt",
2321+
}
2322+
}
2323+
],
2324+
[
2325+
"Done",
2326+
],
2327+
[
2328+
"Sure, I'll create a file for you",
2329+
{
2330+
"tool_use_id": "1",
2331+
"name": "fs_write",
2332+
"args": {
2333+
"command": "create",
2334+
"file_text": "Hello, world!",
2335+
"path": "/file3.txt",
2336+
}
2337+
}
2338+
],
2339+
[
2340+
"Ok I won't.",
2341+
],
2342+
]));
2343+
2344+
ChatContext::new(
2345+
Arc::clone(&ctx),
2346+
Settings::new_fake(),
2347+
std::io::stdout(),
2348+
None,
2349+
InputSource::new_mock(vec![
2350+
"/tools trustall".to_string(),
2351+
"create a new file".to_string(),
2352+
"/tools reset".to_string(),
2353+
"create a new file".to_string(),
2354+
"exit".to_string(),
2355+
]),
2356+
true,
2357+
test_client,
2358+
|| Some(80),
2359+
None,
2360+
load_tools().expect("Tools failed to load."),
2361+
ToolPermissions::new(0),
2362+
)
2363+
.await
2364+
.unwrap()
2365+
.try_chat()
2366+
.await
2367+
.unwrap();
2368+
2369+
assert_eq!(ctx.fs().read_to_string("/file1.txt").await.unwrap(), "Hello, world!\n");
2370+
assert!(!ctx.fs().exists("/file2.txt"));
2371+
}
2372+
20792373
#[test]
20802374
fn test_editor_content_processing() {
20812375
// Since we no longer have template replacement, this test is simplified

0 commit comments

Comments
 (0)