From 383fb21c07b57d75ad7c757b67ba714e151101b7 Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Mon, 2 Feb 2026 22:01:09 -0800 Subject: [PATCH] Add workspace integration tests - Create tests/workspace_integration.rs with 10 database tests - Export MockEmbeddings for use in integration tests - Tests cover: read/write, append, nested paths, delete, memory ops, daily log, FTS search, hybrid search, list_all, system_prompt Requires: DATABASE_URL=postgres://localhost/near_agent_test Co-Authored-By: Claude Opus 4.5 --- src/workspace/embeddings.rs | 7 +- src/workspace/mod.rs | 2 +- tests/workspace_integration.rs | 351 +++++++++++++++++++++++++++++++++ 3 files changed, 356 insertions(+), 4 deletions(-) create mode 100644 tests/workspace_integration.rs diff --git a/src/workspace/embeddings.rs b/src/workspace/embeddings.rs index 5c80d1ed..6cfee5a2 100644 --- a/src/workspace/embeddings.rs +++ b/src/workspace/embeddings.rs @@ -214,19 +214,20 @@ impl EmbeddingProvider for OpenAiEmbeddings { } /// A mock embedding provider for testing. -#[cfg(test)] +/// +/// Generates deterministic embeddings based on text hash. +/// Useful for unit and integration tests. pub struct MockEmbeddings { dimension: usize, } -#[cfg(test)] impl MockEmbeddings { + /// Create a new mock embeddings provider with the given dimension. pub fn new(dimension: usize) -> Self { Self { dimension } } } -#[cfg(test)] #[async_trait] impl EmbeddingProvider for MockEmbeddings { fn dimension(&self) -> usize { diff --git a/src/workspace/mod.rs b/src/workspace/mod.rs index 12a4721c..841b9431 100644 --- a/src/workspace/mod.rs +++ b/src/workspace/mod.rs @@ -48,7 +48,7 @@ mod search; pub use chunker::{ChunkConfig, chunk_document}; pub use document::{MemoryChunk, MemoryDocument, WorkspaceEntry, paths}; -pub use embeddings::{EmbeddingProvider, OpenAiEmbeddings}; +pub use embeddings::{EmbeddingProvider, MockEmbeddings, OpenAiEmbeddings}; pub use repository::Repository; pub use search::{SearchConfig, SearchResult}; diff --git a/tests/workspace_integration.rs b/tests/workspace_integration.rs new file mode 100644 index 00000000..aa05aadf --- /dev/null +++ b/tests/workspace_integration.rs @@ -0,0 +1,351 @@ +//! Integration tests for the workspace module. +//! +//! Requires a running PostgreSQL with pgvector extension. +//! Set DATABASE_URL=postgres://localhost/near_agent_test + +use std::sync::Arc; + +use near_agent::workspace::{MockEmbeddings, SearchConfig, Workspace, paths}; + +fn get_pool() -> deadpool_postgres::Pool { + let database_url = + std::env::var("DATABASE_URL").unwrap_or_else(|_| "postgres://localhost/near_agent_test".to_string()); + + let config: tokio_postgres::Config = database_url.parse().expect("Invalid DATABASE_URL"); + + let mgr = deadpool_postgres::Manager::new(config, tokio_postgres::NoTls); + deadpool_postgres::Pool::builder(mgr) + .max_size(4) + .build() + .expect("Failed to create pool") +} + +async fn cleanup_user(pool: &deadpool_postgres::Pool, user_id: &str) { + let conn = pool.get().await.expect("Failed to get connection"); + conn.execute( + "DELETE FROM memory_documents WHERE user_id = $1", + &[&user_id], + ) + .await + .ok(); +} + +#[tokio::test] +async fn test_workspace_write_and_read() { + let pool = get_pool(); + let user_id = "test_write_read"; + cleanup_user(&pool, user_id).await; + + let workspace = Workspace::new(user_id, pool.clone()); + + // Write a file + let doc = workspace + .write("README.md", "# Hello World\n\nThis is a test.") + .await + .expect("Failed to write"); + + assert_eq!(doc.path, "README.md"); + assert!(doc.content.contains("Hello World")); + + // Read it back + let doc2 = workspace.read("README.md").await.expect("Failed to read"); + assert_eq!(doc2.content, "# Hello World\n\nThis is a test."); + + // Cleanup + cleanup_user(&pool, user_id).await; +} + +#[tokio::test] +async fn test_workspace_append() { + let pool = get_pool(); + let user_id = "test_append"; + cleanup_user(&pool, user_id).await; + + let workspace = Workspace::new(user_id, pool.clone()); + + // Write initial content + workspace + .write("notes.md", "Line 1") + .await + .expect("Failed to write"); + + // Append more + workspace + .append("notes.md", "Line 2") + .await + .expect("Failed to append"); + + // Read and verify + let doc = workspace.read("notes.md").await.expect("Failed to read"); + assert_eq!(doc.content, "Line 1\nLine 2"); + + cleanup_user(&pool, user_id).await; +} + +#[tokio::test] +async fn test_workspace_nested_paths() { + let pool = get_pool(); + let user_id = "test_nested"; + cleanup_user(&pool, user_id).await; + + let workspace = Workspace::new(user_id, pool.clone()); + + // Write nested files + workspace + .write("projects/alpha/README.md", "# Alpha") + .await + .expect("Failed to write alpha"); + workspace + .write("projects/alpha/notes.md", "Notes here") + .await + .expect("Failed to write notes"); + workspace + .write("projects/beta/README.md", "# Beta") + .await + .expect("Failed to write beta"); + + // List root + let root = workspace.list("").await.expect("Failed to list root"); + assert_eq!(root.len(), 1); // just "projects/" + assert!(root[0].is_directory); + assert_eq!(root[0].name(), "projects"); + + // List projects + let projects = workspace + .list("projects") + .await + .expect("Failed to list projects"); + assert_eq!(projects.len(), 2); // alpha/, beta/ + + // List alpha + let alpha = workspace + .list("projects/alpha") + .await + .expect("Failed to list alpha"); + assert_eq!(alpha.len(), 2); // README.md, notes.md + + cleanup_user(&pool, user_id).await; +} + +#[tokio::test] +async fn test_workspace_delete() { + let pool = get_pool(); + let user_id = "test_delete"; + cleanup_user(&pool, user_id).await; + + let workspace = Workspace::new(user_id, pool.clone()); + + // Write and verify exists + workspace + .write("temp.md", "temporary") + .await + .expect("Failed to write"); + assert!(workspace.exists("temp.md").await.expect("exists failed")); + + // Delete + workspace.delete("temp.md").await.expect("Failed to delete"); + + // Verify gone + assert!(!workspace.exists("temp.md").await.expect("exists failed")); + + cleanup_user(&pool, user_id).await; +} + +#[tokio::test] +async fn test_workspace_memory_operations() { + let pool = get_pool(); + let user_id = "test_memory_ops"; + cleanup_user(&pool, user_id).await; + + let workspace = Workspace::new(user_id, pool.clone()); + + // Append to memory + workspace + .append_memory("User prefers dark mode") + .await + .expect("Failed to append memory"); + workspace + .append_memory("User's timezone is PST") + .await + .expect("Failed to append memory"); + + // Read memory + let memory = workspace.memory().await.expect("Failed to get memory"); + assert!(memory.content.contains("dark mode")); + assert!(memory.content.contains("PST")); + // Entries should be separated by double newline + assert!(memory.content.contains("\n\n")); + + cleanup_user(&pool, user_id).await; +} + +#[tokio::test] +async fn test_workspace_daily_log() { + let pool = get_pool(); + let user_id = "test_daily_log"; + cleanup_user(&pool, user_id).await; + + let workspace = Workspace::new(user_id, pool.clone()); + + // Append to daily log (timestamped) + workspace + .append_daily_log("Started working on feature X") + .await + .expect("Failed to append daily log"); + + // Read today's log + let log = workspace.today_log().await.expect("Failed to get today log"); + assert!(log.content.contains("feature X")); + // Should have timestamp prefix like [HH:MM:SS] + assert!(log.content.contains("[")); + + cleanup_user(&pool, user_id).await; +} + +#[tokio::test] +async fn test_workspace_fts_search() { + let pool = get_pool(); + let user_id = "test_fts_search"; + cleanup_user(&pool, user_id).await; + + let workspace = Workspace::new(user_id, pool.clone()); + + // Write some documents + workspace + .write( + "docs/authentication.md", + "# Authentication\n\nThe system uses JWT tokens for authentication.", + ) + .await + .expect("write failed"); + workspace + .write( + "docs/database.md", + "# Database\n\nWe use PostgreSQL with pgvector for vector search.", + ) + .await + .expect("write failed"); + workspace + .write( + "docs/api.md", + "# API\n\nThe REST API uses JSON for request and response bodies.", + ) + .await + .expect("write failed"); + + // Search for JWT (FTS only since no embeddings) + let results = workspace + .search_with_config("JWT authentication", SearchConfig::default().fts_only()) + .await + .expect("search failed"); + + assert!(!results.is_empty(), "Should find results for JWT"); + assert!( + results[0].content.contains("JWT"), + "Top result should contain JWT" + ); + + // Search for PostgreSQL + let results = workspace + .search_with_config("PostgreSQL database", SearchConfig::default().fts_only()) + .await + .expect("search failed"); + + assert!(!results.is_empty(), "Should find results for PostgreSQL"); + assert!( + results[0].content.contains("PostgreSQL"), + "Top result should contain PostgreSQL" + ); + + cleanup_user(&pool, user_id).await; +} + +#[tokio::test] +async fn test_workspace_hybrid_search_with_mock_embeddings() { + let pool = get_pool(); + let user_id = "test_hybrid_search"; + cleanup_user(&pool, user_id).await; + + // Create workspace with mock embeddings (1536 dimensions to match OpenAI) + let embeddings = Arc::new(MockEmbeddings::new(1536)); + let workspace = Workspace::new(user_id, pool.clone()).with_embeddings(embeddings); + + // Write documents + workspace + .write("memory.md", "The user prefers dark mode and vim keybindings.") + .await + .expect("write failed"); + workspace + .write("prefs.md", "Settings: theme=dark, editor=vim, font=monospace") + .await + .expect("write failed"); + + // Hybrid search + let results = workspace + .search("dark theme preference", 5) + .await + .expect("search failed"); + + assert!(!results.is_empty(), "Should find results"); + // At least one result should be a hybrid match (found by both FTS and vector) + // or we should have results from either method + + cleanup_user(&pool, user_id).await; +} + +#[tokio::test] +async fn test_workspace_list_all() { + let pool = get_pool(); + let user_id = "test_list_all"; + cleanup_user(&pool, user_id).await; + + let workspace = Workspace::new(user_id, pool.clone()); + + // Write files at various depths + workspace.write("README.md", "root").await.unwrap(); + workspace.write("docs/intro.md", "intro").await.unwrap(); + workspace.write("docs/api/rest.md", "rest").await.unwrap(); + workspace.write("src/main.md", "main").await.unwrap(); + + // List all + let all = workspace.list_all().await.expect("list_all failed"); + assert_eq!(all.len(), 4); + assert!(all.contains(&"README.md".to_string())); + assert!(all.contains(&"docs/intro.md".to_string())); + assert!(all.contains(&"docs/api/rest.md".to_string())); + assert!(all.contains(&"src/main.md".to_string())); + + cleanup_user(&pool, user_id).await; +} + +#[tokio::test] +async fn test_workspace_system_prompt() { + let pool = get_pool(); + let user_id = "test_system_prompt"; + cleanup_user(&pool, user_id).await; + + let workspace = Workspace::new(user_id, pool.clone()); + + // Write identity files + workspace + .write(paths::AGENTS, "You are a helpful assistant.") + .await + .unwrap(); + workspace + .write(paths::SOUL, "Be kind and thorough.") + .await + .unwrap(); + workspace + .write(paths::USER, "Name: Alice") + .await + .unwrap(); + + // Get system prompt + let prompt = workspace.system_prompt().await.expect("system_prompt failed"); + + assert!(prompt.contains("helpful assistant"), "Should include AGENTS.md"); + assert!(prompt.contains("kind and thorough"), "Should include SOUL.md"); + assert!(prompt.contains("Alice"), "Should include USER.md"); + + cleanup_user(&pool, user_id).await; +}