"""Tests for memory modules.""" import pytest import time from fusionagi.memory.working import WorkingMemory from fusionagi.memory.episodic import EpisodicMemory from fusionagi.memory.reflective import ReflectiveMemory class TestWorkingMemory: """Test WorkingMemory functionality.""" def test_get_set(self): """Test basic get/set operations.""" wm = WorkingMemory() wm.set("session1", "key1", "value1") assert wm.get("session1", "key1") == "value1" assert wm.get("session1", "key2") is None assert wm.get("session1", "key2", "default") == "default" def test_append(self): """Test append to list.""" wm = WorkingMemory() wm.append("s1", "items", "a") wm.append("s1", "items", "b") wm.append("s1", "items", "c") items = wm.get_list("s1", "items") assert items == ["a", "b", "c"] def test_append_converts_non_list(self): """Test append converts non-list values to list.""" wm = WorkingMemory() wm.set("s1", "val", "single") wm.append("s1", "val", "new") items = wm.get_list("s1", "val") assert items == ["single", "new"] def test_has_and_keys(self): """Test has() and keys() methods.""" wm = WorkingMemory() wm.set("s1", "k1", "v1") wm.set("s1", "k2", "v2") assert wm.has("s1", "k1") is True assert wm.has("s1", "k3") is False assert set(wm.keys("s1")) == {"k1", "k2"} def test_delete(self): """Test delete operation.""" wm = WorkingMemory() wm.set("s1", "key", "value") assert wm.has("s1", "key") result = wm.delete("s1", "key") assert result is True assert not wm.has("s1", "key") # Delete non-existent returns False result = wm.delete("s1", "key") assert result is False def test_clear_session(self): """Test clearing a session.""" wm = WorkingMemory() wm.set("s1", "k1", "v1") wm.set("s1", "k2", "v2") wm.set("s2", "k1", "v1") wm.clear_session("s1") assert not wm.session_exists("s1") assert wm.session_exists("s2") def test_context_summary(self): """Test context summary generation.""" wm = WorkingMemory() wm.set("s1", "scalar", "hello") wm.set("s1", "list_val", [1, 2, 3, 4, 5]) wm.set("s1", "dict_val", {"a": 1, "b": 2}) summary = wm.get_context_summary("s1") assert "scalar" in summary assert summary["scalar"] == "hello" assert summary["list_val"]["type"] == "list" assert summary["list_val"]["count"] == 5 assert summary["dict_val"]["type"] == "dict" def test_session_count(self): """Test session counting.""" wm = WorkingMemory() assert wm.session_count() == 0 wm.set("s1", "k", "v") wm.set("s2", "k", "v") assert wm.session_count() == 2 class TestEpisodicMemory: """Test EpisodicMemory functionality.""" def test_append_and_get_by_task(self): """Test appending and retrieving by task.""" em = EpisodicMemory() em.append("task1", {"step": "s1", "result": "ok"}) em.append("task1", {"step": "s2", "result": "ok"}) em.append("task2", {"step": "s1", "result": "fail"}) task1_entries = em.get_by_task("task1") assert len(task1_entries) == 2 assert task1_entries[0]["step"] == "s1" task2_entries = em.get_by_task("task2") assert len(task2_entries) == 1 def test_get_by_type(self): """Test retrieving by event type.""" em = EpisodicMemory() em.append("t1", {"data": 1}, event_type="step_done") em.append("t1", {"data": 2}, event_type="step_done") em.append("t1", {"data": 3}, event_type="step_failed") done_events = em.get_by_type("step_done") assert len(done_events) == 2 failed_events = em.get_by_type("step_failed") assert len(failed_events) == 1 def test_get_recent(self): """Test getting recent entries.""" em = EpisodicMemory() for i in range(10): em.append("task", {"n": i}) recent = em.get_recent(limit=5) assert len(recent) == 5 assert recent[0]["n"] == 5 # 5th entry assert recent[4]["n"] == 9 # 10th entry def test_query_with_filter(self): """Test custom query filter.""" em = EpisodicMemory() em.append("t1", {"score": 0.9, "type": "a"}) em.append("t1", {"score": 0.5, "type": "b"}) em.append("t1", {"score": 0.8, "type": "a"}) high_scores = em.query(lambda e: e.get("score", 0) > 0.7) assert len(high_scores) == 2 def test_task_summary(self): """Test task summary generation.""" em = EpisodicMemory() em.append("task1", {"success": True}, event_type="step_done") em.append("task1", {"success": True}, event_type="step_done") em.append("task1", {"error": "fail"}, event_type="step_failed") summary = em.get_task_summary("task1") assert summary["count"] == 3 assert summary["success_count"] == 2 assert summary["failure_count"] == 1 assert "step_done" in summary["event_types"] def test_statistics(self): """Test overall statistics.""" em = EpisodicMemory() em.append("t1", {}, event_type="type_a") em.append("t2", {}, event_type="type_b") stats = em.get_statistics() assert stats["total_entries"] == 2 assert stats["task_count"] == 2 assert stats["event_type_count"] == 2 def test_clear(self): """Test clearing all entries.""" em = EpisodicMemory() em.append("t1", {}) em.append("t2", {}) em.clear() assert em.get_statistics()["total_entries"] == 0 class TestReflectiveMemory: """Test ReflectiveMemory functionality.""" def test_add_and_get_lessons(self): """Test adding and retrieving lessons.""" rm = ReflectiveMemory() rm.add_lesson({"content": "Don't repeat mistakes", "source": "critic"}) rm.add_lesson({"content": "Plan before acting", "source": "critic"}) lessons = rm.get_lessons() assert len(lessons) == 2 assert lessons[0]["content"] == "Don't repeat mistakes" def test_add_and_get_heuristics(self): """Test adding and retrieving heuristics.""" rm = ReflectiveMemory() rm.set_heuristic("strategy1", "Check dependencies first") rm.set_heuristic("strategy2", "Validate inputs early") heuristics = rm.get_all_heuristics() assert len(heuristics) == 2 assert rm.get_heuristic("strategy1") == "Check dependencies first" def test_get_recent_limits(self): """Test limits on recent retrieval.""" rm = ReflectiveMemory() for i in range(10): rm.add_lesson({"id": i, "content": f"Lesson {i}"}) recent = rm.get_lessons(limit=5) assert len(recent) == 5 # Should get the last 5 assert recent[0]["id"] == 5 assert recent[4]["id"] == 9