mempalace--mempalace/tests/test_migrate.py

"""Tests for destructive-operation safety in mempalace.migrate."""

import errno
import os
import sqlite3
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest

from mempalace.migrate import (
    _restore_stale_palace,
    collection_write_roundtrip_works,
    extract_drawers_from_sqlite,
    migrate,
)


def test_migrate_requires_palace_database(tmp_path, capsys):
    palace_dir = tmp_path / "palace"
    palace_dir.mkdir()

    result = migrate(str(palace_dir))

    out = capsys.readouterr().out
    assert result is False
    assert "No palace database found" in out


def test_migrate_aborts_without_confirmation(tmp_path, capsys):
    palace_dir = tmp_path / "palace"
    palace_dir.mkdir()
    # Presence of chroma.sqlite3 is the safety gate; validity is mocked below.
    (palace_dir / "chroma.sqlite3").write_text("db")

    mock_chromadb = SimpleNamespace(
        __version__="0.6.0",
        PersistentClient=MagicMock(side_effect=Exception("unreadable")),
    )

    with (
        patch.dict("sys.modules", {"chromadb": mock_chromadb}),
        patch("mempalace.migrate.detect_chromadb_version", return_value="0.5.x"),
        patch(
            "mempalace.migrate.extract_drawers_from_sqlite",
            return_value=[{"id": "id1", "document": "doc", "metadata": {"wing": "w", "room": "r"}}],
        ),
        patch("builtins.input", return_value="n"),
        patch("mempalace.migrate.shutil.copytree") as mock_copytree,
        patch("mempalace.migrate.shutil.rmtree") as mock_rmtree,
    ):
        result = migrate(str(palace_dir))

    out = capsys.readouterr().out
    assert result is False
    assert "Aborted." in out
    mock_copytree.assert_not_called()
    mock_rmtree.assert_not_called()


def test_restore_stale_palace_with_clean_destination(tmp_path):
    """Rollback when no partial copy exists at palace_path."""
    palace_path = tmp_path / "palace"
    stale_path = tmp_path / "palace.old"
    stale_path.mkdir()
    (stale_path / "chroma.sqlite3").write_bytes(b"original")

    _restore_stale_palace(str(palace_path), str(stale_path))

    assert palace_path.is_dir()
    assert (palace_path / "chroma.sqlite3").read_bytes() == b"original"
    assert not stale_path.exists()


def test_restore_stale_palace_clears_partial_copy(tmp_path):
    """Rollback must remove a partially-copied palace_path before restoring.

    Simulates the Qodo-reported hazard: shutil.move() began creating
    palace_path, then failed. A bare os.replace(stale, palace_path) would
    trip on the existing destination; _restore_stale_palace must clear it.
    """
    palace_path = tmp_path / "palace"
    stale_path = tmp_path / "palace.old"

    stale_path.mkdir()
    (stale_path / "chroma.sqlite3").write_bytes(b"original")

    palace_path.mkdir()
    (palace_path / "half-copied.bin").write_bytes(b"garbage")

    _restore_stale_palace(str(palace_path), str(stale_path))

    assert palace_path.is_dir()
    assert (palace_path / "chroma.sqlite3").read_bytes() == b"original"
    assert not (palace_path / "half-copied.bin").exists()
    assert not stale_path.exists()


def test_restore_stale_palace_logs_and_swallows_on_failure(tmp_path, capsys):
    """If restore itself fails, log both paths — don't raise from rollback."""
    palace_path = tmp_path / "palace"
    stale_path = tmp_path / "palace.old"
    stale_path.mkdir()

    # Force os.replace to fail deterministically.
    with patch("mempalace.migrate.os.replace", side_effect=OSError("boom")):
        _restore_stale_palace(str(palace_path), str(stale_path))

    out = capsys.readouterr().out
    assert "CRITICAL" in out
    assert os.fspath(palace_path) in out
    assert os.fspath(stale_path) in out


class _FakeGetResult:
    def __init__(self, ids):
        self.ids = ids


class _WritableFakeCollection:
    def __init__(self):
        self.ids = set()
        self.deleted = []

    def upsert(self, *, ids, documents, metadatas):
        self.ids.update(ids)

    def get(self, *, ids, include=None):
        return _FakeGetResult([drawer_id for drawer_id in ids if drawer_id in self.ids])

    def delete(self, *, ids=None, where=None):
        for drawer_id in ids or []:
            self.ids.discard(drawer_id)
            self.deleted.append(drawer_id)


class _SilentWriteDropCollection(_WritableFakeCollection):
    def upsert(self, *, ids, documents, metadatas):
        return None


class _SilentDeleteDropCollection(_WritableFakeCollection):
    def delete(self, *, ids=None, where=None):
        self.deleted.extend(ids or [])


def test_collection_write_roundtrip_works_when_probe_persists_and_deletes():
    col = _WritableFakeCollection()

    assert collection_write_roundtrip_works(col) is True
    assert col.ids == set()
    assert len(col.deleted) == 1


def test_collection_write_roundtrip_fails_when_upsert_silently_drops():
    col = _SilentWriteDropCollection()

    assert collection_write_roundtrip_works(col) is False
    assert col.ids == set()


def test_collection_write_roundtrip_fails_when_delete_silently_drops():
    col = _SilentDeleteDropCollection()

    assert collection_write_roundtrip_works(col) is False
    assert len(col.ids) == 1


def _make_minimal_chromadb_sqlite(tmp_path):
    """Build a SQLite file with the minimal schema extract_drawers_from_sqlite reads."""
    db = tmp_path / "chroma.sqlite3"
    conn = sqlite3.connect(str(db))
    conn.executescript(
        """
        CREATE TABLE embeddings (id INTEGER PRIMARY KEY, embedding_id TEXT);
        CREATE TABLE embedding_metadata (
            id INTEGER, key TEXT,
            string_value TEXT, int_value INTEGER,
            float_value REAL, bool_value INTEGER
        );
        INSERT INTO embeddings VALUES (1, 'd-001');
        INSERT INTO embedding_metadata VALUES (1, 'chroma:document', 'hello', NULL, NULL, NULL);
        INSERT INTO embedding_metadata VALUES (1, 'wing', 'personal', NULL, NULL, NULL);
        INSERT INTO embedding_metadata VALUES (1, 'room', '2026-04-26', NULL, NULL, NULL);
        """
    )
    conn.commit()
    conn.close()
    return str(db)


def test_extract_drawers_returns_drawers(tmp_path):
    db_path = _make_minimal_chromadb_sqlite(tmp_path)
    drawers = extract_drawers_from_sqlite(db_path)
    assert len(drawers) == 1
    assert drawers[0]["id"] == "d-001"
    assert drawers[0]["document"] == "hello"
    assert drawers[0]["metadata"] == {"wing": "personal", "room": "2026-04-26"}


def test_migrate_dry_run_rebuilds_when_collection_is_readable_but_not_writable(tmp_path, capsys):
    palace_dir = tmp_path / "palace"
    palace_dir.mkdir()
    (palace_dir / "chroma.sqlite3").write_text("db")

    fake_col = MagicMock()
    fake_col.count.return_value = 102

    drawers = [
        {
            "id": "id1",
            "document": "hello",
            "metadata": {"wing": "test-wing", "room": "general"},
        }
    ]

    with (
        patch("mempalace.migrate.detect_chromadb_version", return_value="1.x"),
        patch("mempalace.backends.chroma.ChromaBackend") as mock_backend,
        patch(
            "mempalace.migrate.collection_write_roundtrip_works", return_value=False
        ) as mock_probe,
        patch(
            "mempalace.migrate.extract_drawers_from_sqlite", return_value=drawers
        ) as mock_extract,
    ):
        mock_backend.backend_version.return_value = "1.5.8"
        mock_backend.return_value.get_collection.return_value = fake_col

        result = migrate(str(palace_dir), dry_run=True)

    out = capsys.readouterr().out

    assert result is True
    mock_probe.assert_called_once_with(fake_col)
    mock_extract.assert_called_once_with(
        os.path.join(os.path.abspath(os.fspath(palace_dir)), "chroma.sqlite3")
    )

    assert "readable by chromadb 1.5.8, but write/delete verification failed" in out
    assert "Rebuilding from SQLite" in out
    assert "Extracted 1 drawers from SQLite" in out
    assert "DRY RUN" in out


def test_migrate_cleans_temp_palace_on_chromadb_failure(tmp_path):
    """If chromadb fails after the temp palace is created, mkdtemp's
    directory must be removed — without try/finally it leaked into the
    system temp root forever."""
    import tempfile as _tempfile

    palace_dir = tmp_path / "palace"
    palace_dir.mkdir()
    (palace_dir / "chroma.sqlite3").write_text("db")

    captured_temp_paths = []
    real_mkdtemp = _tempfile.mkdtemp

    def tracking_mkdtemp(*args, **kwargs):
        path = real_mkdtemp(*args, **kwargs)
        captured_temp_paths.append(path)
        return path

    failing_backend = MagicMock()
    # First ChromaBackend().get_collection() must raise so we drop into
    # the SQL-extraction path; the second ChromaBackend().get_or_create_collection()
    # raises to trigger the cleanup we are testing.
    failing_backend.get_collection.side_effect = Exception("unreadable")
    failing_backend.get_or_create_collection.side_effect = RuntimeError("chromadb boom")

    import mempalace.backends.chroma as _chroma_mod

    with (
        patch("mempalace.migrate.detect_chromadb_version", return_value="0.5.x"),
        patch(
            "mempalace.migrate.extract_drawers_from_sqlite",
            return_value=[{"id": "id1", "document": "doc", "metadata": {"wing": "w", "room": "r"}}],
        ),
        patch("builtins.input", return_value="y"),
        patch("mempalace.migrate.shutil.copytree"),
        patch("mempalace.migrate.tempfile.mkdtemp", side_effect=tracking_mkdtemp),
        patch.object(_chroma_mod, "ChromaBackend", return_value=failing_backend),
    ):
        try:
            migrate(str(palace_dir), confirm=True)
        except Exception:
            pass

    assert captured_temp_paths, "mkdtemp was never called — flow short-circuited"
    for p in captured_temp_paths:
        assert not os.path.exists(p), f"temp palace was not cleaned up: {p}"


def test_migrate_prunes_old_pre_migrate_backups(tmp_path, monkeypatch):
    """Repeated migrations must not accumulate full-palace copies forever.

    The backup + prune happen right after copytree, before the (mocked)
    chromadb step, so even a migration that fails afterward still trims the
    backup set. We let copytree run for real so the fresh backup exists on
    disk for the prune to evaluate.
    """
    palace_dir = tmp_path / "palace"
    palace_dir.mkdir()
    (palace_dir / "chroma.sqlite3").write_text("db")

    # Pre-seed 3 stale .pre-migrate.* sibling dirs with old mtimes.
    for i in range(3):
        stale = tmp_path / f"palace.pre-migrate.2026010{i}_000000"
        stale.mkdir()
        (stale / "chroma.sqlite3").write_text("old")
        os.utime(stale, (1_700_000_000 + i, 1_700_000_000 + i))

    monkeypatch.setenv("MEMPALACE_MAX_BACKUPS", "2")

    failing_backend = MagicMock()
    failing_backend.get_collection.side_effect = Exception("unreadable")
    failing_backend.get_or_create_collection.side_effect = RuntimeError("chromadb boom")

    import mempalace.backends.chroma as _chroma_mod

    with (
        patch("mempalace.migrate.detect_chromadb_version", return_value="0.5.x"),
        patch(
            "mempalace.migrate.extract_drawers_from_sqlite",
            return_value=[{"id": "id1", "document": "doc", "metadata": {"wing": "w", "room": "r"}}],
        ),
        patch("builtins.input", return_value="y"),
        patch.object(_chroma_mod, "ChromaBackend", return_value=failing_backend),
    ):
        try:
            migrate(str(palace_dir), confirm=True)
        except Exception:
            pass

    backups = sorted(p.name for p in tmp_path.glob("palace.pre-migrate.*"))
    # 3 stale + 1 fresh = 4 created; retention keeps only the 2 newest.
    assert len(backups) == 2
    # The two oldest stale backups must be gone.
    assert "palace.pre-migrate.20260100_000000" not in backups
    assert "palace.pre-migrate.20260101_000000" not in backups


def test_migrate_restores_palace_on_swap_failure(tmp_path, capsys):
    """End-to-end coverage for swap-failure rollback.

    `migrate` swaps the old palace aside via `os.replace` rather than
    deleting it. If `os.replace(temp_palace, palace_path)` raises EXDEV
    (cross-filesystem) AND its `shutil.move` fallback ALSO fails,
    `_restore_stale_palace` rolls back by renaming the aside-copy back
    into place. This exercises that full failure path through the public
    `migrate()` entry point; develop already has unit-level tests for the
    helper itself.
    """
    palace_dir = tmp_path / "palace"
    palace_dir.mkdir()
    (palace_dir / "chroma.sqlite3").write_text("dummy db")
    # Sentinel file we verify survives the failed swap via rename-aside rollback.
    (palace_dir / "sentinel.txt").write_text("original")

    fake_col = MagicMock()
    fake_col.count.return_value = 1
    fake_col.add.return_value = None

    drawers = [{"id": "id1", "document": "doc", "metadata": {"wing": "w", "room": "r"}}]

    # Selective os.replace mock: pass-through for the rename-aside (call A,
    # palace -> palace.old) and the rollback (call C, palace.old -> palace);
    # raise EXDEV exactly once on the swap-in (call B, temp -> palace).
    real_os_replace = os.replace
    fail_state = {"swap_in_failed": False}

    def selective_replace(src, dst):
        if os.fspath(dst) == os.fspath(palace_dir) and not fail_state["swap_in_failed"]:
            fail_state["swap_in_failed"] = True
            raise OSError(errno.EXDEV, "Invalid cross-device link")
        return real_os_replace(src, dst)

    with (
        patch("mempalace.migrate.detect_chromadb_version", return_value="0.5.x"),
        patch("mempalace.backends.chroma.ChromaBackend") as mock_backend_cls,
        patch("mempalace.migrate.collection_write_roundtrip_works", return_value=False),
        patch("mempalace.migrate.extract_drawers_from_sqlite", return_value=drawers),
        patch("mempalace.migrate.confirm_destructive_action", return_value=True),
        patch("mempalace.migrate.os.replace", side_effect=selective_replace),
        patch(
            "mempalace.migrate.shutil.move",
            side_effect=OSError("fallback move also failed"),
        ),
        pytest.raises(OSError),
    ):
        mock_backend_cls.backend_version.return_value = "1.5.4"
        mock_backend_cls.return_value.get_collection.return_value = fake_col
        mock_backend_cls.return_value.get_or_create_collection.return_value = fake_col
        migrate(str(palace_dir))

    # Palace directory restored from the rename-aside copy.
    assert palace_dir.is_dir(), "palace directory missing after rollback"
    sentinel = palace_dir / "sentinel.txt"
    assert sentinel.is_file(), "sentinel file not restored"
    assert sentinel.read_text() == "original", "restored contents differ from original"

    # Pre-migrate backup remains on disk for post-mortem.
    backups = [p for p in tmp_path.iterdir() if p.name.startswith("palace.pre-migrate.")]
    assert backups, "pre-migrate backup directory missing"

    # Stale .old aside-copy was consumed by the rollback (renamed back).
    stale_path = tmp_path / "palace.old"
    assert not stale_path.exists(), "stale .old should have been consumed by rollback"

    # No CRITICAL message — rollback succeeded cleanly.
    out = capsys.readouterr().out
    assert "CRITICAL" not in out