andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-01-25 16:18:13 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Apply automatic linting fixes to tests directory

Browse Source 
Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
Devin AI
2025-05-12 13:31:07 +00:00
parent ad1ea46bbb
commit 46621113af
62 changed files with 1738 additions and 1821 deletions
Download Patch File Download Diff File Expand all files Collapse all files

93
tests/storage/test_mem0_storage.py
View File

@@ -1,34 +1,28 @@
import os
from unittest.mock import MagicMock, patch
import pytest
from mem0.client.main import MemoryClient
from mem0.memory.main import Memory
from crewai.agent import Agent
from crewai.crew import Crew
from crewai.memory.storage.mem0_storage import Mem0Storage
from crewai.task import Task
# Define the class (if not already defined)
class MockCrew:
def __init__(self, memory_config):
def __init__(self, memory_config) -> None:
self.memory_config = memory_config
self.agents = [MagicMock(role="Test Agent")]
@pytest.fixture
def mock_mem0_memory():
"""Fixture to create a mock Memory instance"""
mock_memory = MagicMock(spec=Memory)
return mock_memory
"""Fixture to create a mock Memory instance."""
return MagicMock(spec=Memory)
@pytest.fixture
def mem0_storage_with_mocked_config(mock_mem0_memory):
"""Fixture to create a Mem0Storage instance with mocked dependencies"""
"""Fixture to create a Mem0Storage instance with mocked dependencies."""
# Patch the Memory class to return our mock
with patch("mem0.memory.main.Memory.from_config", return_value=mock_mem0_memory) as mock_from_config:
config = {
@@ -63,15 +57,15 @@ def mem0_storage_with_mocked_config(mock_mem0_memory):
memory_config={
"provider": "mem0",
"config": {"user_id": "test_user", "local_mem0_config": config},
}
},
)
mem0_storage = Mem0Storage(type="short_term", crew=crew)
return mem0_storage, mock_from_config, config
def test_mem0_storage_initialization(mem0_storage_with_mocked_config, mock_mem0_memory):
"""Test that Mem0Storage initializes correctly with the mocked config"""
def test_mem0_storage_initialization(mem0_storage_with_mocked_config, mock_mem0_memory) -> None:
"""Test that Mem0Storage initializes correctly with the mocked config."""
mem0_storage, mock_from_config, config = mem0_storage_with_mocked_config
assert mem0_storage.memory_type == "short_term"
assert mem0_storage.memory is mock_mem0_memory
@@ -80,15 +74,13 @@ def test_mem0_storage_initialization(mem0_storage_with_mocked_config, mock_mem0_
@pytest.fixture
def mock_mem0_memory_client():
"""Fixture to create a mock MemoryClient instance"""
mock_memory = MagicMock(spec=MemoryClient)
return mock_memory
"""Fixture to create a mock MemoryClient instance."""
return MagicMock(spec=MemoryClient)
@pytest.fixture
def mem0_storage_with_memory_client_using_config_from_crew(mock_mem0_memory_client):
"""Fixture to create a Mem0Storage instance with mocked dependencies"""
"""Fixture to create a Mem0Storage instance with mocked dependencies."""
# We need to patch the MemoryClient before it's instantiated
with patch.object(MemoryClient, "__new__", return_value=mock_mem0_memory_client):
crew = MockCrew(
@@ -100,17 +92,15 @@ def mem0_storage_with_memory_client_using_config_from_crew(mock_mem0_memory_clie
"org_id": "my_org_id",
"project_id": "my_project_id",
},
}
},
)
mem0_storage = Mem0Storage(type="short_term", crew=crew)
return mem0_storage
return Mem0Storage(type="short_term", crew=crew)
@pytest.fixture
def mem0_storage_with_memory_client_using_explictly_config(mock_mem0_memory_client, mock_mem0_memory):
"""Fixture to create a Mem0Storage instance with mocked dependencies"""
"""Fixture to create a Mem0Storage instance with mocked dependencies."""
# We need to patch both MemoryClient and Memory to prevent actual initialization
with patch.object(MemoryClient, "__new__", return_value=mock_mem0_memory_client), \
patch.object(Memory, "__new__", return_value=mock_mem0_memory):
@@ -124,19 +114,18 @@ def mem0_storage_with_memory_client_using_explictly_config(mock_mem0_memory_clie
"org_id": "my_org_id",
"project_id": "my_project_id",
},
}
},
)
new_config = {"provider": "mem0", "config": {"api_key": "new-api-key"}}
mem0_storage = Mem0Storage(type="short_term", crew=crew, config=new_config)
return mem0_storage
return Mem0Storage(type="short_term", crew=crew, config=new_config)
def test_mem0_storage_with_memory_client_initialization(
mem0_storage_with_memory_client_using_config_from_crew, mock_mem0_memory_client
):
"""Test Mem0Storage initialization with MemoryClient"""
mem0_storage_with_memory_client_using_config_from_crew, mock_mem0_memory_client,
) -> None:
"""Test Mem0Storage initialization with MemoryClient."""
assert (
mem0_storage_with_memory_client_using_config_from_crew.memory_type
== "short_term"
@@ -149,7 +138,7 @@ def test_mem0_storage_with_memory_client_initialization(
def test_mem0_storage_with_explict_config(
mem0_storage_with_memory_client_using_explictly_config,
):
) -> None:
expected_config = {"provider": "mem0", "config": {"api_key": "new-api-key"}}
assert (
mem0_storage_with_memory_client_using_explictly_config.config == expected_config
@@ -160,17 +149,17 @@ def test_mem0_storage_with_explict_config(
)
def test_save_method_with_memory_oss(mem0_storage_with_mocked_config):
"""Test save method for different memory types"""
def test_save_method_with_memory_oss(mem0_storage_with_mocked_config) -> None:
"""Test save method for different memory types."""
mem0_storage, _, _ = mem0_storage_with_mocked_config
mem0_storage.memory.add = MagicMock()
# Test short_term memory type (already set in fixture)
test_value = "This is a test memory"
test_metadata = {"key": "value"}
mem0_storage.save(test_value, test_metadata)
mem0_storage.memory.add.assert_called_once_with(
test_value,
agent_id="Test_Agent",
@@ -179,28 +168,28 @@ def test_save_method_with_memory_oss(mem0_storage_with_mocked_config):
)
def test_save_method_with_memory_client(mem0_storage_with_memory_client_using_config_from_crew):
"""Test save method for different memory types"""
def test_save_method_with_memory_client(mem0_storage_with_memory_client_using_config_from_crew) -> None:
"""Test save method for different memory types."""
mem0_storage = mem0_storage_with_memory_client_using_config_from_crew
mem0_storage.memory.add = MagicMock()
# Test short_term memory type (already set in fixture)
test_value = "This is a test memory"
test_metadata = {"key": "value"}
mem0_storage.save(test_value, test_metadata)
mem0_storage.memory.add.assert_called_once_with(
test_value,
agent_id="Test_Agent",
infer=False,
metadata={"type": "short_term", "key": "value"},
output_format="v1.1"
output_format="v1.1",
)
def test_search_method_with_memory_oss(mem0_storage_with_mocked_config):
"""Test search method for different memory types"""
def test_search_method_with_memory_oss(mem0_storage_with_mocked_config) -> None:
"""Test search method for different memory types."""
mem0_storage, _, _ = mem0_storage_with_mocked_config
mock_results = {"results": [{"score": 0.9, "content": "Result 1"}, {"score": 0.4, "content": "Result 2"}]}
mem0_storage.memory.search = MagicMock(return_value=mock_results)
@@ -208,18 +197,18 @@ def test_search_method_with_memory_oss(mem0_storage_with_mocked_config):
results = mem0_storage.search("test query", limit=5, score_threshold=0.5)
mem0_storage.memory.search.assert_called_once_with(
query="test query",
limit=5,
query="test query",
limit=5,
agent_id="Test_Agent",
user_id="test_user"
user_id="test_user",
)
assert len(results) == 1
assert results[0]["content"] == "Result 1"
def test_search_method_with_memory_client(mem0_storage_with_memory_client_using_config_from_crew):
"""Test search method for different memory types"""
def test_search_method_with_memory_client(mem0_storage_with_memory_client_using_config_from_crew) -> None:
"""Test search method for different memory types."""
mem0_storage = mem0_storage_with_memory_client_using_config_from_crew
mock_results = {"results": [{"score": 0.9, "content": "Result 1"}, {"score": 0.4, "content": "Result 2"}]}
mem0_storage.memory.search = MagicMock(return_value=mock_results)
@@ -227,12 +216,12 @@ def test_search_method_with_memory_client(mem0_storage_with_memory_client_using_
results = mem0_storage.search("test query", limit=5, score_threshold=0.5)
mem0_storage.memory.search.assert_called_once_with(
query="test query",
limit=5,
agent_id="Test_Agent",
query="test query",
limit=5,
agent_id="Test_Agent",
metadata={"type": "short_term"},
user_id="test_user",
output_format='v1.1'
output_format="v1.1",
)
assert len(results) == 1