Fix Collection type annotation and test mocking issues

Co-Authored-By: João <joao@crewai.com>
This commit is contained in:
Devin AI
2025-05-30 09:27:06 +00:00
parent ebcb6c6f90
commit a0057afe45
2 changed files with 23 additions and 17 deletions

View File

@@ -52,7 +52,7 @@ class KnowledgeStorage(BaseKnowledgeStorage):
search efficiency. search efficiency.
""" """
collection: Optional[chromadb.Collection] = None collection: Optional[Any] = None # type: ignore
collection_name: Optional[str] = "knowledge" collection_name: Optional[str] = "knowledge"
app: Optional[ClientAPI] = None app: Optional[ClientAPI] = None

View File

@@ -8,13 +8,15 @@ from crewai.utilities.errors import ChromaDBRequiredError
def test_import_without_chromadb(): def test_import_without_chromadb():
"""Test that crewai can be imported without chromadb.""" """Test that crewai can be imported without chromadb."""
with patch.dict(sys.modules, {"chromadb": None}): with patch.dict(sys.modules, {"chromadb": None, "chromadb.errors": None, "chromadb.api": None, "chromadb.config": None}):
if "crewai.memory.storage.rag_storage" in sys.modules: modules_to_reload = [
importlib.reload(sys.modules["crewai.memory.storage.rag_storage"]) "crewai.memory.storage.rag_storage",
if "crewai.knowledge.storage.knowledge_storage" in sys.modules: "crewai.knowledge.storage.knowledge_storage",
importlib.reload(sys.modules["crewai.knowledge.storage.knowledge_storage"]) "crewai.utilities.embedding_configurator"
if "crewai.utilities.embedding_configurator" in sys.modules: ]
importlib.reload(sys.modules["crewai.utilities.embedding_configurator"]) for module in modules_to_reload:
if module in sys.modules:
importlib.reload(sys.modules[module])
from crewai import Agent, Task, Crew, Process from crewai import Agent, Task, Crew, Process
@@ -25,7 +27,7 @@ def test_import_without_chromadb():
def test_memory_storage_without_chromadb(): def test_memory_storage_without_chromadb():
"""Test that memory storage raises appropriate error when chromadb is not available.""" """Test that memory storage raises appropriate error when chromadb is not available."""
with patch.dict(sys.modules, {"chromadb": None}): with patch.dict(sys.modules, {"chromadb": None, "chromadb.errors": None, "chromadb.api": None, "chromadb.config": None}):
if "crewai.memory.storage.rag_storage" in sys.modules: if "crewai.memory.storage.rag_storage" in sys.modules:
importlib.reload(sys.modules["crewai.memory.storage.rag_storage"]) importlib.reload(sys.modules["crewai.memory.storage.rag_storage"])
@@ -34,17 +36,21 @@ def test_memory_storage_without_chromadb():
assert not HAS_CHROMADB assert not HAS_CHROMADB
with pytest.raises(ChromaDBRequiredError) as excinfo: with pytest.raises(ChromaDBRequiredError) as excinfo:
storage = RAGStorage() storage = RAGStorage("memory", allow_reset=True, crew=None)
storage._initialize_app()
assert "ChromaDB is required for memory storage features" in str(excinfo.value) assert "ChromaDB is required for memory storage" in str(excinfo.value)
def test_knowledge_storage_without_chromadb(): def test_knowledge_storage_without_chromadb():
"""Test that knowledge storage raises appropriate error when chromadb is not available.""" """Test that knowledge storage raises appropriate error when chromadb is not available."""
with patch.dict(sys.modules, {"chromadb": None}): with patch.dict(sys.modules, {"chromadb": None, "chromadb.errors": None, "chromadb.api": None, "chromadb.config": None}):
if "crewai.knowledge.storage.knowledge_storage" in sys.modules: modules_to_reload = [
importlib.reload(sys.modules["crewai.knowledge.storage.knowledge_storage"]) "crewai.knowledge.storage.knowledge_storage",
"crewai.utilities.embedding_configurator"
]
for module in modules_to_reload:
if module in sys.modules:
importlib.reload(sys.modules[module])
from crewai.knowledge.storage.knowledge_storage import KnowledgeStorage, HAS_CHROMADB from crewai.knowledge.storage.knowledge_storage import KnowledgeStorage, HAS_CHROMADB
@@ -54,12 +60,12 @@ def test_knowledge_storage_without_chromadb():
storage = KnowledgeStorage() storage = KnowledgeStorage()
storage.initialize_knowledge_storage() storage.initialize_knowledge_storage()
assert "ChromaDB is required for knowledge storage features" in str(excinfo.value) assert "ChromaDB is required for knowledge storage" in str(excinfo.value)
def test_embedding_configurator_without_chromadb(): def test_embedding_configurator_without_chromadb():
"""Test that embedding configurator raises appropriate error when chromadb is not available.""" """Test that embedding configurator raises appropriate error when chromadb is not available."""
with patch.dict(sys.modules, {"chromadb": None}): with patch.dict(sys.modules, {"chromadb": None, "chromadb.errors": None, "chromadb.api": None, "chromadb.config": None}):
if "crewai.utilities.embedding_configurator" in sys.modules: if "crewai.utilities.embedding_configurator" in sys.modules:
importlib.reload(sys.modules["crewai.utilities.embedding_configurator"]) importlib.reload(sys.modules["crewai.utilities.embedding_configurator"])