diff --git a/tests/memory/test_memory_reset.py b/tests/memory/test_memory_reset.py index 964bd664b..330405a06 100644 --- a/tests/memory/test_memory_reset.py +++ b/tests/memory/test_memory_reset.py @@ -1,22 +1,34 @@ import os +import tempfile import pytest from crewai.memory import ShortTermMemory, LongTermMemory, EntityMemory from crewai.utilities import EmbeddingConfigurator -def test_memory_reset_with_ollama(): +@pytest.fixture +def temp_db_dir(): + with tempfile.TemporaryDirectory() as tmpdir: + yield tmpdir + +def test_memory_reset_with_ollama(temp_db_dir): os.environ["CREWAI_EMBEDDING_PROVIDER"] = "ollama" os.environ["CREWAI_EMBEDDING_MODEL"] = "llama2" - # Test each memory type - memories = [ShortTermMemory(), LongTermMemory(), EntityMemory()] + memories = [ + ShortTermMemory(path=temp_db_dir), + LongTermMemory(path=temp_db_dir), + EntityMemory(path=temp_db_dir) + ] for memory in memories: - memory.reset() # Should not raise any OpenAI-related errors + memory.reset() -def test_memory_reset_with_openai(): +def test_memory_reset_with_openai(temp_db_dir): os.environ["CREWAI_EMBEDDING_PROVIDER"] = "openai" os.environ["CREWAI_EMBEDDING_MODEL"] = "text-embedding-3-small" - # Test each memory type - memories = [ShortTermMemory(), LongTermMemory(), EntityMemory()] + memories = [ + ShortTermMemory(path=temp_db_dir), + LongTermMemory(path=temp_db_dir), + EntityMemory(path=temp_db_dir) + ] for memory in memories: - memory.reset() # Should work with OpenAI key + memory.reset()