Initial release: iai-mcp v0.1.0

Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: XNLLLLH <XNLLLLH@users.noreply.github.com>

tests/test_migrate_crypto_recover_prior_key.py

@@ -0,0 +1,101 @@
+"""Tests for migrate_crypto_recover_prior_key (prior AES key re-encrypt + swap)."""
+
+from __future__ import annotations
+
+import os
+import secrets
+from datetime import datetime, timezone
+from pathlib import Path
+from uuid import uuid4
+
+import pytest
+from cryptography.exceptions import InvalidTag
+
+from iai_mcp.migrate import migrate_crypto_recover_prior_key
+from iai_mcp.store import MemoryStore
+from iai_mcp.types import MemoryRecord, SCHEMA_VERSION_CURRENT
+
+
+def _minimal_record(literal: str) -> MemoryRecord:
+    rid = uuid4()
+    now = datetime.now(timezone.utc)
+    return MemoryRecord(
+        id=rid,
+        tier="episodic",
+        literal_surface=literal,
+        aaak_index="",
+        embedding=[0.01] * 384,
+        structure_hv=b"\x00" * 1250,
+        community_id=None,
+        centrality=0.0,
+        detail_level=1,
+        pinned=False,
+        stability=0.0,
+        difficulty=0.0,
+        last_reviewed=None,
+        never_decay=False,
+        never_merge=False,
+        provenance=[],
+        created_at=now,
+        updated_at=now,
+        tags=[],
+        language="en",
+        s5_trust_score=0.5,
+        profile_modulation_gain={},
+        schema_version=SCHEMA_VERSION_CURRENT,
+    )
+
+
+def test_recover_prior_key_atomic_swap_and_idempotent(tmp_path: Path) -> None:
+    root = tmp_path / "store"
+    root.mkdir()
+    key_a = secrets.token_bytes(32)
+    key_b = secrets.token_bytes(32)
+    kpath = root / ".crypto.key"
+    kpath.write_bytes(key_a)
+    os.chmod(kpath, 0o600)
+
+    store_a = MemoryStore(path=root, user_id="default")
+    rec = _minimal_record("verbatim-prior-key-recover")
+    store_a.insert(rec)
+    rid = rec.id
+    del store_a
+
+    kpath.write_bytes(key_b)
+    os.chmod(kpath, 0o600)
+    store_b = MemoryStore(path=root, user_id="default")
+    with pytest.raises(InvalidTag):
+        store_b.get(rid)
+
+    out = migrate_crypto_recover_prior_key(store_b, key_a, dry_run=False)
+    assert out.get("no_op") is False
+    assert out.get("records_staged") == 1
+    assert out.get("rows_needed_prior_key") == 1
+
+    got = store_b.get(rid)
+    assert got is not None
+    assert got.literal_surface == "verbatim-prior-key-recover"
+
+    out2 = migrate_crypto_recover_prior_key(store_b, key_a, dry_run=False)
+    assert out2.get("no_op") is True
+    assert out2.get("reason") == "all_rows_decrypt_with_current_key"
+
+
+def test_recover_prior_key_dry_run_counts(tmp_path: Path) -> None:
+    root = tmp_path / "store2"
+    root.mkdir()
+    key_a = secrets.token_bytes(32)
+    key_b = secrets.token_bytes(32)
+    kpath = root / ".crypto.key"
+    kpath.write_bytes(key_a)
+    os.chmod(kpath, 0o600)
+    store_a = MemoryStore(path=root, user_id="default")
+    store_a.insert(_minimal_record("dry-run-count"))
+    del store_a
+    kpath.write_bytes(key_b)
+    os.chmod(kpath, 0o600)
+    store_b = MemoryStore(path=root, user_id="default")
+    out = migrate_crypto_recover_prior_key(store_b, key_a, dry_run=True)
+    assert out.get("dry_run") is True
+    assert out.get("would_stage") == 1
+    assert out.get("rows_needing_prior_key") == 1