From 0dfa5b05d374780210be66263638ddc81b718835 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 9 Feb 2025 23:28:23 +0000 Subject: [PATCH] test: Use custom embedder for cleanup test Co-Authored-By: Joe Moura --- tests/memory/test_memory_reset.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/memory/test_memory_reset.py b/tests/memory/test_memory_reset.py index 7faedc4d9..a8d5dd3ea 100644 --- a/tests/memory/test_memory_reset.py +++ b/tests/memory/test_memory_reset.py @@ -65,8 +65,16 @@ def test_memory_reset_with_missing_api_key(temp_db_dir): def test_memory_reset_cleans_up_files(temp_db_dir): """Test that memory reset properly cleans up database files.""" - os.environ["OPENAI_API_KEY"] = "test-key" - memory = ShortTermMemory(path=str(temp_db_dir)) + class TestEmbedder(EmbeddingFunction): + def __call__(self, input: Documents) -> Embeddings: + if isinstance(input, str): + input = [input] + return [[0.5] * 10] * len(input) + + memory = ShortTermMemory( + path=str(temp_db_dir), + embedder_config={"provider": TestEmbedder()} + ) memory.save("test memory", {"test": "metadata"}) assert any(temp_db_dir.iterdir()) # Directory should have files memory.reset()