from __future__ import annotations
from unittest.mock import MagicMock, call
import pytest
from engram_client.integrations.crewai import (
EngramEntityMemory,
EngramLongTermMemory,
EngramShortTermMemory,
)
@pytest.fixture
def mock_client() -> MagicMock:
client = MagicMock()
client.create_daily.return_value = {"id": 1, "content": "[key] value"}
client.create.return_value = {"id": 2, "content": "[key] value"}
client.search.return_value = {"memories": []}
client.create_identity.return_value = {"canonical_id": "person:alice", "display_name": "Alice"}
client.resolve_identity.return_value = {"canonical_id": "person:alice"}
client.delete.return_value = {"deleted": True}
return client
class TestEngramShortTermMemory:
def test_save_calls_create_daily(self, mock_client: MagicMock) -> None:
stm = EngramShortTermMemory(mock_client, workspace="crewai-stm", ttl_seconds=3600)
stm.save("task_result", "All steps completed")
mock_client.create_daily.assert_called_once_with(
content="[task_result] All steps completed",
tags=["crewai", "short-term", "key:task_result"],
workspace="crewai-stm",
ttl_seconds=3600,
metadata={},
)
def test_save_passes_custom_metadata(self, mock_client: MagicMock) -> None:
stm = EngramShortTermMemory(mock_client)
meta = {"agent": "researcher", "step": 3}
stm.save("intermediate", "partial findings", metadata=meta)
_, kwargs = mock_client.create_daily.call_args
assert kwargs["metadata"] == meta
def test_search_calls_client_search(self, mock_client: MagicMock) -> None:
stm = EngramShortTermMemory(mock_client, workspace="crewai-stm")
stm.search("completed tasks", limit=3)
mock_client.search.assert_called_once_with(
"completed tasks", workspace="crewai-stm", limit=3
)
def test_reset_deletes_found_memories(self, mock_client: MagicMock) -> None:
mock_client.search.return_value = {
"memories": [{"id": 10}, {"id": 11}, {"id": 12}]
}
stm = EngramShortTermMemory(mock_client, workspace="crewai-stm")
stm.reset()
mock_client.search.assert_called_once_with(
"crewai", workspace="crewai-stm", limit=100
)
assert mock_client.delete.call_count == 3
mock_client.delete.assert_any_call(10)
mock_client.delete.assert_any_call(11)
mock_client.delete.assert_any_call(12)
def test_reset_skips_entries_without_id(self, mock_client: MagicMock) -> None:
mock_client.search.return_value = {
"memories": [{"content": "no id here"}, {"id": 5}]
}
stm = EngramShortTermMemory(mock_client)
stm.reset()
mock_client.delete.assert_called_once_with(5)
def test_default_workspace_and_ttl(self, mock_client: MagicMock) -> None:
stm = EngramShortTermMemory(mock_client)
assert stm.workspace == "crewai-stm"
assert stm.ttl_seconds == 3600
class TestEngramLongTermMemory:
def test_save_calls_create(self, mock_client: MagicMock) -> None:
ltm = EngramLongTermMemory(mock_client, crew_name="research-crew")
ltm.save("finding_1", "Paris is the capital of France")
mock_client.create.assert_called_once_with(
content="[finding_1] Paris is the capital of France",
tags=["crewai", "long-term", "key:finding_1"],
workspace="crewai-research-crew",
metadata={},
)
def test_uses_crew_specific_workspace(self, mock_client: MagicMock) -> None:
ltm = EngramLongTermMemory(mock_client, crew_name="sales-team")
assert ltm.workspace == "crewai-sales-team"
def test_custom_workspace_overrides_crew_name(self, mock_client: MagicMock) -> None:
ltm = EngramLongTermMemory(
mock_client, crew_name="ignored", workspace="my-custom-ws"
)
assert ltm.workspace == "my-custom-ws"
def test_search_calls_client_search(self, mock_client: MagicMock) -> None:
ltm = EngramLongTermMemory(mock_client, crew_name="analysis")
ltm.search("capital cities", limit=7)
mock_client.search.assert_called_once_with(
"capital cities", workspace="crewai-analysis", limit=7
)
def test_save_passes_metadata(self, mock_client: MagicMock) -> None:
ltm = EngramLongTermMemory(mock_client)
meta = {"source": "wikipedia", "confidence": 0.95}
ltm.save("fact", "The sky is blue", metadata=meta)
_, kwargs = mock_client.create.call_args
assert kwargs["metadata"] == meta
class TestEngramEntityMemory:
def test_save_entity_calls_create_identity(self, mock_client: MagicMock) -> None:
em = EngramEntityMemory(mock_client)
em.save_entity("Alice Smith", "person", "Lead researcher")
mock_client.create_identity.assert_called_once_with(
canonical_id="person:alice_smith",
display_name="Alice Smith (Lead researcher)",
aliases=[],
)
def test_save_entity_passes_aliases(self, mock_client: MagicMock) -> None:
em = EngramEntityMemory(mock_client)
aliases = ["alice@example.com", "+5511999999999"]
em.save_entity("Alice", "person", "Engineer", aliases=aliases)
_, kwargs = mock_client.create_identity.call_args
assert kwargs["aliases"] == aliases
def test_get_entity_resolves_canonical_id(self, mock_client: MagicMock) -> None:
em = EngramEntityMemory(mock_client)
em.get_entity("Bob Jones", entity_type="person")
mock_client.resolve_identity.assert_called_once_with("person:bob_jones")
def test_get_entity_default_type_is_person(self, mock_client: MagicMock) -> None:
em = EngramEntityMemory(mock_client)
em.get_entity("Carol")
mock_client.resolve_identity.assert_called_once_with("person:carol")
def test_search_entities_calls_client_search(self, mock_client: MagicMock) -> None:
em = EngramEntityMemory(mock_client, workspace="crewai-entities")
em.search_entities("researcher", limit=3)
mock_client.search.assert_called_once_with(
"researcher", workspace="crewai-entities", limit=3
)
def test_canonical_id_normalises_spaces(self, mock_client: MagicMock) -> None:
em = EngramEntityMemory(mock_client)
em.save_entity("Open AI", "company", "AI lab")
_, kwargs = mock_client.create_identity.call_args
assert kwargs["canonical_id"] == "company:open_ai"