update to use TestBase instead

Author: brandonskiser

crates/agent/src/agent/mod.rs

@@ -2171,15 +2171,11 @@ pub enum HookStage {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::util::test::{
-        TestDir,
-        TestProvider,
-    };
+    use crate::util::test::TestBase;
 
     #[tokio::test]
     async fn test_collect_resources() {
-        let mut test_dir = TestDir::new();
-        let test_provider = TestProvider::new_with_base(test_dir.path());
+        let mut test_base = TestBase::new().await;
 
         let files = [
             (".amazonq/rules/first.md", "first"),
@@ -2188,10 +2184,10 @@ mod tests {
         ];
 
         for file in files {
-            test_dir = test_dir.with_file_sys(file, &test_provider).await;
+            test_base = test_base.with_file(file).await;
         }
 
-        let resources = collect_resources(["file://.amazonq/rules/**/*.md", "file://~/home.txt"], &test_provider).await;
+        let resources = collect_resources(["file://.amazonq/rules/**/*.md", "file://~/home.txt"], &test_base).await;
 
         for file in files {
             assert!(resources.iter().any(|r| r.content == file.1));

crates/agent/src/agent/tools/fs_read.rs

@@ -196,11 +196,7 @@ pub struct FileReadContext {}
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::agent::util::test::TestDir;
-    use crate::util::test::{
-        TestBase,
-        TestProvider,
-    };
+    use crate::util::test::TestBase;
 
     #[tokio::test]
     async fn test_fs_read_single_file() {
@@ -227,56 +223,56 @@ mod tests {
 
     #[tokio::test]
     async fn test_fs_read_with_offset_and_limit() {
-        let test_provider = TestProvider::new();
-        let test_dir = TestDir::new()
-            .with_file_sys(("test.txt", "line1\nline2\nline3\nline4\nline5"), &test_provider)
+        let test_base = TestBase::new()
+            .await
+            .with_file(("test.txt", "line1\nline2\nline3\nline4\nline5"))
             .await;
 
         let tool = FsRead {
             ops: vec![FsReadOp {
-                path: test_dir.join("test.txt").to_string_lossy().to_string(),
+                path: test_base.join("test.txt").to_string_lossy().to_string(),
                 limit: Some(2),
                 offset: Some(1),
             }],
         };
 
-        let result = tool.execute(&test_provider).await.unwrap();
+        let result = tool.execute(&test_base).await.unwrap();
         if let ToolExecutionOutputItem::Text(content) = &result.items[0] {
             assert_eq!(content, "line2\nline3");
         }
     }
 
     #[tokio::test]
     async fn test_fs_read_multiple_files() {
-        let test_provider = TestProvider::new();
-        let test_dir = TestDir::new()
-            .with_file_sys(("file1.txt", "content1"), &test_provider)
+        let test_base = TestBase::new()
+            .await
+            .with_file(("file1.txt", "content1"))
             .await
-            .with_file_sys(("file2.txt", "content2"), &test_provider)
+            .with_file(("file2.txt", "content2"))
             .await;
 
         let tool = FsRead {
             ops: vec![
                 FsReadOp {
-                    path: test_dir.join("file1.txt").to_string_lossy().to_string(),
+                    path: test_base.join("file1.txt").to_string_lossy().to_string(),
                     limit: None,
                     offset: None,
                 },
                 FsReadOp {
-                    path: test_dir.join("file2.txt").to_string_lossy().to_string(),
+                    path: test_base.join("file2.txt").to_string_lossy().to_string(),
                     limit: None,
                     offset: None,
                 },
             ],
         };
 
-        let result = tool.execute(&test_provider).await.unwrap();
+        let result = tool.execute(&test_base).await.unwrap();
         assert_eq!(result.items.len(), 2);
     }
 
     #[tokio::test]
     async fn test_fs_read_validate_nonexistent_file() {
-        let test_provider = TestProvider::new();
+        let test_base = TestBase::new().await;
         let tool = FsRead {
             ops: vec![FsReadOp {
                 path: "/nonexistent/file.txt".to_string(),
@@ -285,22 +281,21 @@ mod tests {
             }],
         };
 
-        assert!(tool.validate(&test_provider).await.is_err());
+        assert!(tool.validate(&test_base).await.is_err());
     }
 
     #[tokio::test]
     async fn test_fs_read_validate_directory_path() {
-        let test_provider = TestProvider::new();
-        let test_dir = TestDir::new();
+        let test_base = TestBase::new().await;
 
         let tool = FsRead {
             ops: vec![FsReadOp {
-                path: test_dir.join("").to_string_lossy().to_string(),
+                path: test_base.join("").to_string_lossy().to_string(),
                 limit: None,
                 offset: None,
             }],
         };
 
-        assert!(tool.validate(&test_provider).await.is_err());
+        assert!(tool.validate(&test_base).await.is_err());
     }
 }

crates/agent/src/agent/tools/fs_write.rs

@@ -354,148 +354,155 @@ impl FileLineTracker {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::agent::util::test::TestDir;
-    use crate::util::test::TestProvider;
+    use crate::util::test::TestBase;
 
     #[tokio::test]
     async fn test_create_file() {
-        let test_provider = TestProvider::new();
-        let test_dir = TestDir::new();
+        let test_base = TestBase::new().await;
         let tool = FsWrite::Create(FileCreate {
-            path: test_dir.join("new.txt").to_string_lossy().to_string(),
+            path: test_base.join("new.txt").to_string_lossy().to_string(),
             content: "hello world".to_string(),
         });
 
-        assert!(tool.validate(&test_provider).await.is_ok());
-        assert!(tool.execute(None, &test_provider).await.is_ok());
+        assert!(tool.validate(&test_base).await.is_ok());
+        assert!(tool.execute(None, &test_base).await.is_ok());
 
-        let content = tokio::fs::read_to_string(test_dir.join("new.txt")).await.unwrap();
+        let content = tokio::fs::read_to_string(test_base.join("new.txt")).await.unwrap();
         assert_eq!(content, "hello world");
     }
 
     #[tokio::test]
     async fn test_create_file_with_parent_dirs() {
-        let test_provider = TestProvider::new();
-        let test_dir = TestDir::new();
+        let test_base = TestBase::new().await;
         let tool = FsWrite::Create(FileCreate {
-            path: test_dir.join("nested/dir/file.txt").to_string_lossy().to_string(),
+            path: test_base.join("nested/dir/file.txt").to_string_lossy().to_string(),
             content: "nested content".to_string(),
         });
 
-        assert!(tool.execute(None, &test_provider).await.is_ok());
+        assert!(tool.execute(None, &test_base).await.is_ok());
 
-        let content = tokio::fs::read_to_string(test_dir.join("nested/dir/file.txt"))
+        let content = tokio::fs::read_to_string(test_base.join("nested/dir/file.txt"))
             .await
             .unwrap();
         assert_eq!(content, "nested content");
     }
 
     #[tokio::test]
     async fn test_str_replace_single_occurrence() {
-        let test_provider = TestProvider::new();
-        let test_dir = TestDir::new().with_file_sys(("test.txt", "hello world"), &test_provider).await;
+        let test_base = TestBase::new()
+            .await
+            .with_file(("test.txt", "hello world"))
+            .await;
 
         let tool = FsWrite::StrReplace(StrReplace {
-            path: test_dir.join("test.txt").to_string_lossy().to_string(),
+            path: test_base.join("test.txt").to_string_lossy().to_string(),
             old_str: "world".to_string(),
             new_str: "rust".to_string(),
             replace_all: false,
         });
 
-        assert!(tool.execute(None, &test_provider).await.is_ok());
+        assert!(tool.execute(None, &test_base).await.is_ok());
 
-        let content = tokio::fs::read_to_string(test_dir.join("test.txt")).await.unwrap();
+        let content = tokio::fs::read_to_string(test_base.join("test.txt")).await.unwrap();
         assert_eq!(content, "hello rust");
     }
 
     #[tokio::test]
     async fn test_str_replace_multiple_occurrences() {
-        let test_provider = TestProvider::new();
-        let test_dir = TestDir::new().with_file_sys(("test.txt", "foo bar foo"), &test_provider).await;
+        let test_base = TestBase::new()
+            .await
+            .with_file(("test.txt", "foo bar foo"))
+            .await;
 
         let tool = FsWrite::StrReplace(StrReplace {
-            path: test_dir.join("test.txt").to_string_lossy().to_string(),
+            path: test_base.join("test.txt").to_string_lossy().to_string(),
             old_str: "foo".to_string(),
             new_str: "baz".to_string(),
             replace_all: true,
         });
 
-        assert!(tool.execute(None, &test_provider).await.is_ok());
+        assert!(tool.execute(None, &test_base).await.is_ok());
 
-        let content = tokio::fs::read_to_string(test_dir.join("test.txt")).await.unwrap();
+        let content = tokio::fs::read_to_string(test_base.join("test.txt")).await.unwrap();
         assert_eq!(content, "baz bar baz");
     }
 
     #[tokio::test]
     async fn test_str_replace_no_match() {
-        let test_provider = TestProvider::new();
-        let test_dir = TestDir::new().with_file_sys(("test.txt", "hello world"), &test_provider).await;
+        let test_base = TestBase::new()
+            .await
+            .with_file(("test.txt", "hello world"))
+            .await;
 
         let tool = FsWrite::StrReplace(StrReplace {
-            path: test_dir.join("test.txt").to_string_lossy().to_string(),
+            path: test_base.join("test.txt").to_string_lossy().to_string(),
             old_str: "missing".to_string(),
             new_str: "replacement".to_string(),
             replace_all: false,
         });
 
-        assert!(tool.execute(None, &test_provider).await.is_err());
+        assert!(tool.execute(None, &test_base).await.is_err());
     }
 
     #[tokio::test]
     async fn test_insert_at_line() {
-        let test_provider = TestProvider::new();
-        let test_dir = TestDir::new().with_file_sys(("test.txt", "line1\nline2\nline3"), &test_provider).await;
+        let test_base = TestBase::new()
+            .await
+            .with_file(("test.txt", "line1\nline2\nline3"))
+            .await;
 
         let tool = FsWrite::Insert(Insert {
-            path: test_dir.join("test.txt").to_string_lossy().to_string(),
+            path: test_base.join("test.txt").to_string_lossy().to_string(),
             content: "inserted".to_string(),
             insert_line: Some(1),
         });
 
-        assert!(tool.execute(None, &test_provider).await.is_ok());
+        assert!(tool.execute(None, &test_base).await.is_ok());
 
-        let content = tokio::fs::read_to_string(test_dir.join("test.txt")).await.unwrap();
+        let content = tokio::fs::read_to_string(test_base.join("test.txt")).await.unwrap();
         assert_eq!(content, "line1\ninserted\nline2\nline3");
     }
 
     #[tokio::test]
     async fn test_insert_append() {
-        let test_provider = TestProvider::new();
-        let test_dir = TestDir::new().with_file_sys(("test.txt", "existing"), &test_provider).await;
+        let test_base = TestBase::new()
+            .await
+            .with_file(("test.txt", "existing"))
+            .await;
 
         let tool = FsWrite::Insert(Insert {
-            path: test_dir.join("test.txt").to_string_lossy().to_string(),
+            path: test_base.join("test.txt").to_string_lossy().to_string(),
             content: "appended".to_string(),
             insert_line: None,
         });
 
-        assert!(tool.execute(None, &test_provider).await.is_ok());
+        assert!(tool.execute(None, &test_base).await.is_ok());
 
-        let content = tokio::fs::read_to_string(test_dir.join("test.txt")).await.unwrap();
+        let content = tokio::fs::read_to_string(test_base.join("test.txt")).await.unwrap();
         assert_eq!(content, "existing\nappended");
     }
 
     #[tokio::test]
     async fn test_fs_write_validate_empty_path() {
-        let test_provider = TestProvider::new();
+        let test_base = TestBase::new().await;
         let tool = FsWrite::Create(FileCreate {
             path: "".to_string(),
             content: "content".to_string(),
         });
 
-        assert!(tool.validate(&test_provider).await.is_err());
+        assert!(tool.validate(&test_base).await.is_err());
     }
 
     #[tokio::test]
     async fn test_fs_write_validate_nonexistent_file_for_replace() {
-        let test_provider = TestProvider::new();
+        let test_base = TestBase::new().await;
         let tool = FsWrite::StrReplace(StrReplace {
             path: "/nonexistent/file.txt".to_string(),
             old_str: "old".to_string(),
             new_str: "new".to_string(),
             replace_all: false,
         });
 
-        assert!(tool.validate(&test_provider).await.is_err());
+        assert!(tool.validate(&test_base).await.is_err());
     }
 }