1use 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}