Skip to main content

agentzero_testkit/
lib.rs

1//! Test utilities for AgentZero.
2//!
3//! Provides fake implementations of `Provider`, `MemoryStore`, and `Tool`
4//! for use in unit and integration tests across the workspace.
5
6use agentzero_core::{
7    ChatResult, MemoryEntry, MemoryStore, Provider, Tool, ToolContext, ToolResult,
8};
9use async_trait::async_trait;
10use std::sync::{Arc, Mutex};
11
12#[derive(Default, Clone)]
13pub struct TestMemoryStore {
14    entries: Arc<Mutex<Vec<MemoryEntry>>>,
15}
16
17#[async_trait]
18impl MemoryStore for TestMemoryStore {
19    async fn append(&self, entry: MemoryEntry) -> anyhow::Result<()> {
20        self.entries
21            .lock()
22            .expect("test memory lock poisoned")
23            .push(entry);
24        Ok(())
25    }
26
27    async fn recent(&self, limit: usize) -> anyhow::Result<Vec<MemoryEntry>> {
28        let entries = self.entries.lock().expect("test memory lock poisoned");
29        Ok(entries.iter().rev().take(limit).cloned().collect())
30    }
31}
32
33impl TestMemoryStore {
34    pub fn len(&self) -> usize {
35        self.entries
36            .lock()
37            .expect("test memory lock poisoned")
38            .len()
39    }
40
41    pub fn is_empty(&self) -> bool {
42        self.len() == 0
43    }
44}
45
46#[derive(Clone)]
47pub struct StaticProvider {
48    pub output_text: String,
49}
50
51#[async_trait]
52impl Provider for StaticProvider {
53    async fn complete(&self, _prompt: &str) -> anyhow::Result<ChatResult> {
54        Ok(ChatResult {
55            output_text: self.output_text.clone(),
56            ..Default::default()
57        })
58    }
59}
60
61pub struct FailingProvider;
62
63#[async_trait]
64impl Provider for FailingProvider {
65    async fn complete(&self, _prompt: &str) -> anyhow::Result<ChatResult> {
66        Err(anyhow::anyhow!("testkit provider failure"))
67    }
68}
69
70pub struct EchoTool;
71
72#[async_trait]
73impl Tool for EchoTool {
74    fn name(&self) -> &'static str {
75        "echo"
76    }
77
78    async fn execute(&self, input: &str, _ctx: &ToolContext) -> anyhow::Result<ToolResult> {
79        Ok(ToolResult {
80            output: format!("echoed:{input}"),
81        })
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::{FailingProvider, StaticProvider, TestMemoryStore};
88    use agentzero_core::{MemoryEntry, MemoryStore, Provider};
89
90    #[tokio::test]
91    async fn test_memory_store_round_trip_success() {
92        let store = TestMemoryStore::default();
93        MemoryStore::append(
94            &store,
95            MemoryEntry {
96                role: "user".to_string(),
97                content: "hello".to_string(),
98                ..Default::default()
99            },
100        )
101        .await
102        .expect("append should succeed");
103        assert_eq!(store.len(), 1);
104    }
105
106    #[tokio::test]
107    async fn failing_provider_returns_expected_error() {
108        let provider = FailingProvider;
109        let err = provider
110            .complete("hello")
111            .await
112            .expect_err("provider should fail");
113        assert!(err.to_string().contains("testkit provider failure"));
114    }
115
116    #[tokio::test]
117    async fn static_provider_returns_fixed_output() {
118        let provider = StaticProvider {
119            output_text: "ok".to_string(),
120        };
121        let result = provider
122            .complete("ignored")
123            .await
124            .expect("provider should succeed");
125        assert_eq!(result.output_text, "ok");
126    }
127}