mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 16:48:30 +00:00
test: Use temporary directory for memory reset tests
Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user