a-memory 1.7.0__py3-none-any.whl

tests/shared/test_saga_schema.py

@@ -0,0 +1,48 @@
+from shared.saga.schema import SagaStatus, SagaStepState, SagaState
+import time
+
+
+def test_saga_status_enum():
+    assert SagaStatus.PENDING == "pending"
+    assert SagaStatus.STUCK == "stuck"
+
+
+def test_saga_step_state_validation():
+    step = SagaStepState(name="test_step", status=SagaStatus.PENDING, result={"foo": "bar"})
+    assert step.name == "test_step"
+    assert step.status == SagaStatus.PENDING
+    assert step.result == {"foo": "bar"}
+    assert step.error is None
+
+
+def test_saga_state_validation():
+    saga_id = "abc12345"
+    state = SagaState(
+        saga_id=saga_id,
+        name="test_saga",
+        context={"user_id": 1},
+        steps=[SagaStepState(name="step1", status=SagaStatus.COMPLETED, result={"ok": True})],
+        current_step_index=0,
+        status=SagaStatus.RUNNING,
+        started_at=time.time(),
+    )
+    assert state.saga_id == saga_id
+    assert len(state.steps) == 1
+    assert state.status == SagaStatus.RUNNING
+
+
+def test_saga_state_serialization():
+    state = SagaState(saga_id="deadbeef", name="test", context={}, steps=[], current_step_index=0, status=SagaStatus.PENDING, started_at=time.time())
+    json_data = state.model_dump_json()
+    assert "deadbeef" in json_data
+
+    restored = SagaState.model_validate_json(json_data)
+    assert restored.saga_id == "deadbeef"
+
+
+if __name__ == "__main__":
+    test_saga_status_enum()
+    test_saga_step_state_validation()
+    test_saga_state_validation()
+    test_saga_state_serialization()
+    print("All tests passed!")

tests/test_auth_backup.py

@@ -0,0 +1,84 @@
+"""
+Tests for auth — unique tests only.
+"""
+
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_api_key_create():
+    from features.auth import APIKeyAuth
+
+    auth = APIKeyAuth()
+    key = auth.create_key("alice", "test key")
+    assert key.startswith("ak_")
+    assert len(key) > 20
+
+
+@pytest.mark.asyncio
+async def test_api_key_verify():
+    from features.auth import APIKeyAuth
+
+    auth = APIKeyAuth()
+    key = auth.create_key("alice", "test key")
+    info = auth.verify(key)
+    assert info is not None
+    assert info["user_id"] == "alice"
+    assert info["label"] == "test key"
+
+
+@pytest.mark.asyncio
+async def test_api_key_revoke():
+    from features.auth import APIKeyAuth
+
+    auth = APIKeyAuth()
+    key = auth.create_key("alice", "test key")
+    assert auth.verify(key) is not None
+    revoked = auth.revoke(key)
+    assert revoked is True
+    assert auth.verify(key) is None
+
+
+@pytest.mark.asyncio
+async def test_api_key_list():
+    from features.auth import APIKeyAuth
+
+    auth = APIKeyAuth()
+    auth.create_key("alice", "key1")
+    auth.create_key("alice", "key2")
+    keys = auth.list_keys()
+    assert len(keys) >= 2
+
+
+@pytest.mark.asyncio
+async def test_bearer_auth():
+    from features.auth import BearerAuth
+
+    ba = BearerAuth()
+    token = ba.get_token()
+    assert token.startswith("mt_")
+    assert ba.verify("Bearer " + token) is True
+    assert ba.verify("Bearer invalid") is False
+    assert ba.verify("") is False
+
+
+@pytest.mark.asyncio
+async def test_bearer_rotate():
+    from features.auth import BearerAuth
+
+    ba = BearerAuth()
+    old_token = ba.get_token()
+    new_token = ba.rotate()
+    assert old_token != new_token
+    assert ba.verify("Bearer " + new_token) is True
+
+
+@pytest.mark.asyncio
+async def test_mcp_lifespan():
+    from mcp_server.server import lifespan, mcp
+
+    async with lifespan(mcp) as ctx:
+        assert ctx is not None
+        assert hasattr(ctx, "mm")
+        assert hasattr(ctx, "user_wiki")
+        assert hasattr(ctx, "agent_wiki")

tests/test_auth_crypto.py

@@ -0,0 +1,87 @@
+"""End-to-end tests for auth with encryption."""
+
+import json
+import os
+
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def restore_master_key():
+    """Restore MCP_MASTER_KEY after each test to prevent cross-test contamination."""
+    original = os.environ.get("MCP_MASTER_KEY")
+    yield
+    if original is not None:
+        os.environ["MCP_MASTER_KEY"] = original
+    else:
+        os.environ.pop("MCP_MASTER_KEY", None)
+    from features import secrets
+
+    secrets._master_cache.clear()
+
+
+def test_legacy_plain_json_gets_rotated(tmp_path):
+    """If file is plain JSON, it should be auto-encrypted on load."""
+    keys_file = tmp_path / "api_keys.json"
+    keys_file.write_text(json.dumps({"ak_old": {"user_id": "u1", "label": "old", "enabled": True, "created_at": 1.0}}))
+    from features.auth import APIKeyAuth
+
+    auth = APIKeyAuth(keys_file=str(keys_file))
+    keys = auth.list_keys()
+    assert len(keys) == 1
+    assert keys[0]["user_id"] == "u1"
+
+    # Verify data survived the rotation — read with fresh instance
+    auth2 = APIKeyAuth(keys_file=str(keys_file))
+    keys2 = auth2.list_keys()
+    assert len(keys2) == 1
+    assert keys2[0]["user_id"] == "u1"
+
+
+def test_create_then_verify_round_trip(tmp_path):
+    keys_file = tmp_path / "api_keys.json"
+    from features.auth import APIKeyAuth
+
+    auth = APIKeyAuth(keys_file=str(keys_file))
+    key = auth.create_key("alice", "test-key")
+    out = auth.verify(key)
+    assert out == {"user_id": "alice", "label": "test-key"}
+
+
+def test_wrong_master_key_fails(tmp_path):
+    keys_file = tmp_path / "api_keys.json"
+    from features.auth import APIKeyAuth
+
+    auth = APIKeyAuth(keys_file=str(keys_file))
+    auth.create_key("bob", "test-label")
+
+    # Change master key
+    os.environ["MCP_MASTER_KEY"] = "different-secret"
+    from features import secrets
+
+    secrets._master_cache.clear()
+
+    # Should fail to decrypt (returns empty dict, not exception)
+    auth2 = APIKeyAuth(keys_file=str(keys_file))
+    assert auth2.list_keys() == []  # cannot read old keys
+
+
+def test_bearer_token_encrypted(tmp_path):
+    token_file = tmp_path / "bearer_token.json"
+    from features.auth import BearerAuth
+
+    auth = BearerAuth(token_file=str(token_file))
+    token = auth.get_token()
+    assert token.startswith("mt_")
+
+    # Verify round-trip works
+    assert auth.verify(f"Bearer {token}") is True
+    assert auth.verify("Bearer invalid") is False
+
+    # Rotation produces a new token
+    old_token = token
+    new_token = auth.rotate()
+    assert new_token.startswith("mt_")
+    assert new_token != old_token
+    assert auth.verify(f"Bearer {new_token}") is True
+    assert auth.verify(f"Bearer {old_token}") is False

tests/test_conflict_bm25.py

@@ -0,0 +1,89 @@
+"""Tests for B3: BM25 + char-trigram conflict similarity."""
+
+import pytest
+
+from rag.conflict import (
+    ConflictResolver,
+    bm25_pair_similarity,
+    char_ngram_jaccard,
+    smart_similarity,
+)
+
+
+def test_jaccard_basic():
+    assert char_ngram_jaccard("hello", "hello") == 1.0
+    assert char_ngram_jaccard("abc", "xyz") == 0.0
+
+
+def test_bm25_pair_basic():
+    s = bm25_pair_similarity("redis cluster pipelining", "redis cluster issue")
+    assert 0.0 < s < 1.0
+    s_unrelated = bm25_pair_similarity("redis cluster", "cats and dogs playing")
+    assert s > s_unrelated
+
+
+def test_smart_similarity_short_uses_ngrams():
+    a = "Python is best"
+    b = "Python best"
+    s = smart_similarity(a, b)
+    char_ngram_jaccard(a, b, 3)
+    assert s > 0.5
+
+
+def test_smart_similarity_long_uses_both_signals():
+    a = "Python multi-paragraph text about multi-threading and memory model"
+    b = "Python multi-paragraph text about concurrency and memory model"
+    s = smart_similarity(a, b)
+    assert s > 0.4
+
+
+def test_smart_similarity_rejects_unrelated():
+    s = smart_similarity("Python is the best language", "Cats love fish in rivers")
+    assert s < 0.15
+
+
+def test_smart_similarity_empty_inputs():
+    assert smart_similarity("", "abc") == 0.0
+    assert smart_similarity("abc", "") == 0.0
+    assert smart_similarity("", "") == 0.0
+
+
+@pytest.mark.asyncio
+async def test_resolve_archives_and_audits(tmp_path):
+    from shared.connection import AsyncConnectionManager
+
+    cm = AsyncConnectionManager(base_dir=str(tmp_path))
+    cr = ConflictResolver(cm=cm)
+    await cr._init_db()
+
+    # Create conflict group
+    conn = await cm.get("memory.db")
+    # Init audit table for logging
+    from features.audit_trail import AuditTrail
+
+    at = AuditTrail(cm=cm)
+    await at._init_db()
+
+    await conn.execute(
+        "INSERT INTO memory_conflicts (user_id, content, is_conflict, conflict_group_id) VALUES (?, ?, 1, ?)", ("u", "old version", "grp1")
+    )
+    await conn.execute(
+        "INSERT INTO memory_conflicts (user_id, content, is_conflict, conflict_group_id) VALUES (?, ?, 1, ?)", ("u", "new version", "grp1")
+    )
+    await conn.commit()
+
+    # Resolve: keep_id=2
+    result = await cr.resolve("grp1", keep_id=2)
+    assert result is True
+
+    # Verify only kept entry remains (conflict_group_id set to NULL after resolve)
+    cur = await conn.execute("SELECT id, content FROM memory_conflicts WHERE user_id='u'")
+    rows = await cur.fetchall()
+    assert len(rows) == 1
+    assert rows[0][1] == "new version"
+
+    # Verify audit log
+    cur = await conn.execute("SELECT action, details FROM audit_log WHERE action='conflict_resolved'")
+    audit_row = await cur.fetchone()
+    assert audit_row is not None
+    assert "conflict_group_id" in str(audit_row[1])

tests/test_core/test_core.py

@@ -0,0 +1,11 @@
+"""Tests for core/ module — unique tests only."""
+
+from core.reflex import ReflexBuffer
+
+
+def test_reflex_buffer():
+    buf = ReflexBuffer(max_size=5)
+    buf.add(role="user", content="Hello", tokens=5)
+    buf.add(role="assistant", content="Hi", tokens=3)
+    assert buf.size() == 2
+    assert buf.get_recent(1)[0].content == "Hi"

tests/test_db_maintenance.py

@@ -0,0 +1,74 @@
+"""features.db_maintenance: size thresholds + auto-VACUUM on fragmentation."""
+
+import sqlite3
+
+import pytest
+
+from features.db_maintenance import run_db_maintenance
+from shared.connection import AsyncConnectionManager
+
+
+@pytest.fixture
+async def cm(tmp_path, monkeypatch):
+    manager = AsyncConnectionManager(base_dir=str(tmp_path))
+    # maintenance reads connection_manager.base_dir for the file listing —
+    # point it at this test's tmp and restore afterwards.
+    from shared import connection as conn_mod
+
+    orig_dir = conn_mod.connection_manager.base_dir
+    orig_conns = dict(conn_mod.connection_manager._conns)
+    conn_mod.connection_manager.base_dir = tmp_path
+    conn_mod.connection_manager._conns.clear()
+    yield manager
+    conn_mod.connection_manager.base_dir = orig_dir
+    conn_mod.connection_manager._conns.clear()
+    conn_mod.connection_manager._conns.update(orig_conns)
+
+
+async def _make_fragmented(path, rows: int) -> None:
+    """Create a DB with real dead pages: bulk insert then delete."""
+    conn = sqlite3.connect(str(path))
+    conn.execute("CREATE TABLE t (v TEXT)")
+    conn.executemany("INSERT INTO t VALUES (?)", [("x" * 2000,)] * rows)
+    conn.commit()
+    conn.execute("DELETE FROM t")
+    conn.commit()
+    page_count = conn.execute("PRAGMA page_count").fetchone()[0]
+    freelist = conn.execute("PRAGMA freelist_count").fetchone()[0]
+    assert freelist / max(page_count, 1) > 0.25, "fixture must be fragmented"
+    conn.close()
+
+
+@pytest.mark.asyncio
+async def test_vacuum_reclaims_fragmented_db(cm, tmp_path):
+    db = tmp_path / "memory.db"
+    await _make_fragmented(db, rows=3000)
+
+    # Direct probe first: the vacuum primitive itself must reclaim space.
+    from features.db_maintenance import _vacuum_if_fragmented
+
+    freed, did = await _vacuum_if_fragmented(cm, db, min_mb=1, ratio_threshold=0.25)
+    assert did and freed > 0, f"direct vacuum: did={did} freed={freed}"
+
+    reports = await run_db_maintenance(cm=cm, vacuum_min_mb=50)
+    mem = [r for r in reports if r.name == "memory.db"]
+    assert mem, "memory.db must be reported"
+
+    conn = sqlite3.connect(str(db))
+    assert conn.execute("PRAGMA integrity_check").fetchone()[0] == "ok"
+    assert conn.execute("PRAGMA freelist_count").fetchone()[0] == 0
+
+
+@pytest.mark.asyncio
+async def test_thresholds_reported(cm):
+    """Small clean DB -> ok level; report carries all data-dir files."""
+    import sqlite3 as sq
+
+    conn = sq.connect(str(cm.base_dir / "projects.db"))
+    conn.execute("CREATE TABLE t (v TEXT)")
+    conn.commit()
+    conn.close()
+    reports = await run_db_maintenance(cm=cm)
+    names = {r.name for r in reports}
+    assert "projects.db" in names
+    assert all(r.level in ("ok", "warn", "alert") for r in reports)

tests/test_dream_buffer_flow.py

@@ -0,0 +1,69 @@
+"""DreamBuffer restore: hook stages output, sweep consolidates into L4."""
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from core.projects import ProjectMemory  # noqa: F401 — ensures package import parity
+from shared.connection import AsyncConnectionManager
+from shared.dream_buffer import DreamBuffer
+from shared.migrations import MigrationManager
+
+
+@pytest.fixture
+async def cm(tmp_path):
+    manager = AsyncConnectionManager(base_dir=tmp_path)
+    await MigrationManager(cm=manager).migrate()
+    return manager
+
+
+def _make_ctx(app):
+    ctx = MagicMock()
+    ctx.request_context.lifespan_context = app
+    return ctx
+
+
+@pytest.mark.asyncio
+async def test_dream_hook_stages_into_buffer(tmp_path):
+    """The user-layer dream_buffer handler stages summaries via the caller's cm."""
+    from hooks.user_hooks import UserHooks
+
+    cm = AsyncConnectionManager(base_dir=tmp_path)
+
+    hooks = UserHooks(user_id="du")
+    mem = MagicMock()
+    mem._cm = cm
+    res = await hooks._dream_buffer({"summary": "deployed v2 to prod", "user_id": "du"}, mem=mem)
+    assert res["action"] == "add_to_staging"
+
+    buf = DreamBuffer(cm=cm, layer="user")
+    items = await buf.get_staging("du")
+    assert len(items) == 1 and "v2" in items[0]["content"], f"staged={items}"
+
+    # agent layer stays empty — staging is layer-isolated
+    agent_buf = DreamBuffer(cm=cm, layer="agent")
+    assert await agent_buf.get_staging("du") == []
+
+
+@pytest.mark.asyncio
+async def test_sweep_consolidates_staging_into_l4(cm):
+    """consolidate_staging promotes staged content; clear empties the buffer."""
+    from lifecycle.consolidation import ConsolidationEngine
+
+    buf = DreamBuffer(cm=cm, layer="user")
+    await buf.add("su", session_id="dream", content="critical decision: switched to wal mode always", importance=0.9)
+
+    items = await buf.get_staging("su")
+    engine = ConsolidationEngine(cm=cm, layer="user")
+    res = await engine.consolidate_staging("su", items, min_importance=0.7)
+    assert res["promoted"] == 1
+
+    await buf.clear_staging("su")
+    assert await buf.get_staging("su") == []
+
+    # landed in L4 as a fact
+    from core.memory import CoreMemory
+
+    l4 = CoreMemory(cm=cm, layer="user")
+    hits = await l4.search("su", "wal mode", limit=5)
+    assert any("wal mode" in h["value"] for h in hits)

tests/test_features/test_backup_path_safety.py

@@ -0,0 +1,31 @@
+"""Tests for backup path traversal prevention."""
+
+import asyncio
+import json
+
+import pytest
+
+from features.backup import BackupManager
+
+
+@pytest.fixture
+def bm(tmp_path):
+    data_dir = tmp_path / "data"
+    data_dir.mkdir()
+    return BackupManager(base_dir=str(data_dir))
+
+
+@pytest.mark.parametrize(
+    ("manifest_files", "should_reject"),
+    [
+        (["../../etc/crontab", "memory.db"], True),
+        (["/etc/passwd"], True),
+    ],
+)
+def test_restore_rejects_traversal(bm, manifest_files, should_reject):
+    backup_dir = bm.backup_dir / "crafted"
+    backup_dir.mkdir()
+    manifest = {"files": manifest_files, "created_at": "2026-01-01T00:00:00"}
+    (backup_dir / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8")
+    with pytest.raises(ValueError, match="escapes base directory"):
+        asyncio.run(bm.restore("crafted"))

tests/test_features/test_features.py

@@ -0,0 +1,14 @@
+"""Tests for features/ module — unique tests only."""
+
+import asyncio
+
+
+def test_compression():
+    from features.compression import MemoryCompressor
+
+    async def t():
+        mc = MemoryCompressor()
+        stats = await mc.get_stats("test_feat")
+        assert "core_memory" in stats
+
+    asyncio.run(t())

tests/test_features/test_import_export_path_safety.py

@@ -0,0 +1,32 @@
+"""Tests for import_export path traversal prevention."""
+
+import pytest
+
+from features.import_export import ImportExport
+
+
+@pytest.fixture
+def ie(tmp_path):
+    export_dir = tmp_path / "exports"
+    export_dir.mkdir()
+
+    class FakeCM:
+        def __init__(self, base):
+            self._base = base
+
+        @property
+        def base_dir(self):
+            return self._base
+
+    obj = ImportExport.__new__(ImportExport)
+    obj._cm = FakeCM(tmp_path)
+    obj.export_dir = export_dir
+    return obj
+
+
+@pytest.mark.parametrize("path", ["../../etc/passwd", "/etc/passwd"])
+def test_import_rejects_traversal(ie, path):
+    import asyncio
+
+    with pytest.raises(ValueError, match="escapes base directory"):
+        asyncio.run(ie.import_user(path))