From 8f5610f6ff42abfead40ebfc99c610906a71c834 Mon Sep 17 00:00:00 2001
From: jcarver989 <117696+jcarver989@users.noreply.github.com>
Date: Thu, 30 Jul 2026 05:15:00 +0000
Subject: [PATCH] chore: scheduled code-cleanup

---
 crates/aether-core/tests/agent/agent_tests.rs | 42 ++++---------------
 crates/aether-core/tests/agent/retry_tests.rs | 38 +++++------------
 2 files changed, 18 insertions(+), 62 deletions(-)

diff --git a/crates/aether-core/tests/agent/agent_tests.rs b/crates/aether-core/tests/agent/agent_tests.rs
index 2aed4cd9..53070326 100644
--- a/crates/aether-core/tests/agent/agent_tests.rs
+++ b/crates/aether-core/tests/agent/agent_tests.rs
@@ -5,7 +5,8 @@ use std::time::Duration;
 use aether_core::{
     events::{AgentEvent, Command, TurnOutcome, UserCommand},
     testing::{
-        agent_event, test_agent, {AddNumbersRequest, AddNumbersResult, DivideNumbersRequest, SlowToolRequest},
+        AddNumbersRequest, AddNumbersResult, DivideNumbersRequest, SlowToolRequest, agent_event, content_events,
+        test_agent,
     },
 };
 use llm::testing::{FakeLlmProvider, llm_response};
@@ -17,25 +18,6 @@ fn split_json_in_half(input: &str) -> (&str, &str) {
     input.split_at(split)
 }
 
-/// Strips turn/call lifecycle noise, leaving only the content events the
-/// `agent_event` builder describes.
-fn content_events(events: Vec<AgentEvent>) -> Vec<AgentEvent> {
-    events
-        .into_iter()
-        .filter(|event| {
-            !matches!(
-                event,
-                AgentEvent::Turn(
-                    TurnEvent::Started { .. }
-                        | TurnEvent::RetryScheduled { .. }
-                        | TurnEvent::LlmCallStarted { .. }
-                        | TurnEvent::LlmCallEnded { .. }
-                ) | AgentEvent::Tool(ToolEvent::ExecutionStarted { .. } | ToolEvent::DefinitionsUpdated { .. })
-            )
-        })
-        .collect()
-}
-
 #[tokio::test]
 async fn test_text_message() -> Result<(), Box<dyn Error>> {
     let id = "message_1";
@@ -52,20 +34,12 @@ async fn test_text_message() -> Result<(), Box<dyn Error>> {
 #[tokio::test]
 async fn test_llm_call_lifecycle_reports_model_and_usage() -> Result<(), Box<dyn Error>> {
     let model: llm::LlmModel = "codex:gpt-5.5".parse()?;
-    let llm =
-        FakeLlmProvider::new(vec![llm_response("msg_1").text(&["hi"]).usage(120, 7).build()]).with_model(model.clone());
-    let (tx, mut rx, _handle) = aether_core::core::agent(llm).spawn().await?;
-
-    tx.send(Command::UserCommand(UserCommand::Text { content: vec![llm::ContentBlock::text("hello")] })).await?;
-
-    let mut events = Vec::new();
-    while let Some(event) = rx.recv().await {
-        let is_turn_end = event.turn_outcome().is_some();
-        events.push(event);
-        if is_turn_end {
-            break;
-        }
-    }
+    let events = test_agent()
+        .model(model.clone())
+        .llm_responses(&[llm_response("msg_1").text(&["hi"]).usage(120, 7).build()])
+        .user_text("hello")
+        .run()
+        .await?;
 
     let started = events
         .iter()
diff --git a/crates/aether-core/tests/agent/retry_tests.rs b/crates/aether-core/tests/agent/retry_tests.rs
index ba84165e..431c5e03 100644
--- a/crates/aether-core/tests/agent/retry_tests.rs
+++ b/crates/aether-core/tests/agent/retry_tests.rs
@@ -3,7 +3,7 @@ use std::error::Error;
 use std::time::Duration;
 
 use aether_core::core::RetryConfig;
-use aether_core::events::{AgentEvent, Command, TurnOutcome, UserCommand};
+use aether_core::events::{AgentEvent, TurnOutcome};
 use aether_core::testing::test_agent;
 use llm::{LlmError, LlmResponse};
 
@@ -195,42 +195,24 @@ async fn rate_limited_error_is_retried() -> Result<(), Box<dyn Error>> {
 
 #[tokio::test(start_paused = true)]
 async fn cancel_during_retry_wait_aborts_pending_retry() -> Result<(), Box<dyn Error>> {
-    use aether_core::core::agent;
-    use llm::testing::FakeLlmProvider;
+    use aether_core::testing::TestScenario;
 
     let attempts: Vec<Vec<Result<LlmResponse, LlmError>>> = vec![
         vec![Err(LlmError::ServerError { status: Some(503), message: "boom".into() })],
         vec![Ok(LlmResponse::start("msg_2")), Ok(LlmResponse::text("should not see this")), Ok(LlmResponse::done())],
     ];
 
-    let llm = FakeLlmProvider::from_results(attempts);
-    let captured = llm.captured_contexts();
-
     // Long retry delay; with virtual time it never elapses unless we advance.
     let retry = RetryConfig { max_attempts: 5, base_delay: Duration::from_mins(1), max_delay: Duration::from_mins(1) };
 
-    let (tx, mut rx, _handle) = agent(llm).retry(retry).spawn().await?;
-
-    tx.send(Command::UserCommand(UserCommand::Text { content: vec![llm::ContentBlock::text("go")] })).await?;
-
-    loop {
-        match rx.recv().await {
-            Some(AgentEvent::Turn(TurnEvent::RetryScheduled { attempt: 1, .. })) => break,
-            Some(_) => {}
-            None => panic!("channel closed before the retry was scheduled"),
-        }
-    }
-
-    tx.send(Command::UserCommand(UserCommand::Cancel)).await?;
+    let result = test_agent()
+        .retry_config(retry)
+        .llm_result_responses(&attempts)
+        .scenario(TestScenario::new().user_text("go").wait_for_retry(1).cancel().wait_for_turn_end())
+        .run_with_context()
+        .await?;
 
-    let mut messages = Vec::new();
-    while let Some(msg) = rx.recv().await {
-        let is_turn_end = matches!(msg, AgentEvent::Turn(TurnEvent::Ended { .. }));
-        messages.push(msg);
-        if is_turn_end {
-            break;
-        }
-    }
+    let messages = &result.messages;
 
     let retry_started = messages
         .iter()
@@ -241,7 +223,7 @@ async fn cancel_during_retry_wait_aborts_pending_retry() -> Result<(), Box<dyn E
     assert!(has_cancelled, "expected the turn to end as cancelled, got {messages:?}");
 
     // The retry should never have fired — only the original failed call counts.
-    let captured = captured.lock().unwrap();
+    let captured = result.captured_contexts.lock().unwrap();
     assert_eq!(captured.len(), 1, "retry must not fire after cancel; expected 1 LLM call");
 
     Ok(())