FabledGrimoire/tests/test_integration_backup_rollback.py

"""Real-Postgres proof that a full backup rolls this instance back by id.

The lane shares one database, so every test here takes its backup of the
whole store, changes its own rows, and rolls back to that backup: what it
puts back is exactly what was there when it started.
"""
from datetime import datetime, timezone
from unittest.mock import AsyncMock, patch

import pytest
import pytest_asyncio
from sqlalchemy import func, select

from grimoire.models import async_session
from grimoire.models.note import Note
from grimoire.models.project import Project
from grimoire.models.user import User
from grimoire.restore import run
from grimoire.services import backup, backup_job
from grimoire.services.api_keys import create_api_key, lookup_key
from grimoire.services.backup_file import write_archive
from grimoire.services.instance import adopt_instance_id, get_instance_id
from tests.helpers import ensure_user

pytestmark = [pytest.mark.integration, pytest.mark.usefixtures("_dispose_engine")]

OWNER = "rollback_owner"


@pytest_asyncio.fixture
async def seeded():
    """One user with a project and two tasks, committed, and a full backup
    taken right after."""
    async with async_session() as s:
        user = await ensure_user(s, OWNER)
        # A clean slate for this user: an earlier test's rows would otherwise
        # be in this test's backup too.
        for model in (Note, Project):
            for row in (await s.execute(select(model).where(model.user_id == user.id))).scalars():
                await s.delete(row)
        await s.flush()
        project = Project(user_id=user.id, title="Rollback drill")
        s.add(project)
        await s.flush()
        kept = Note(user_id=user.id, project_id=project.id, title="before",
                    body="as backed up", note_type="note", task_kind="work", status="todo")
        doomed = Note(user_id=user.id, project_id=project.id, title="deleted later",
                      body="", note_type="note", task_kind="work", status="todo")
        s.add_all([kept, doomed])
        await s.commit()
        ids = {"user": user.id, "project": project.id, "kept": kept.id, "doomed": doomed.id}
    key, _ = await create_api_key(ids["user"], "rollback drill", "read")
    payload = await backup.export_full_backup()
    return ids, payload, key


def _own(payload: dict, table: str, user_id: int) -> list[dict]:
    return sorted(
        (row for row in payload[table] if row.get("user_id") == user_id),
        key=lambda row: row["id"],
    )


async def _change_everything(ids: dict) -> int:
    """Edit one task, delete the other, add a task and a project. Returns the
    new task's id."""
    async with async_session() as s:
        kept = await s.get(Note, ids["kept"])
        kept.title = "after"
        kept.body = "edited since the backup"
        await s.delete(await s.get(Note, ids["doomed"]))
        extra = Project(user_id=ids["user"], title="made after the backup")
        s.add(extra)
        await s.flush()
        new = Note(user_id=ids["user"], project_id=extra.id, title="new since",
                   body="", note_type="note", task_kind="work", status="todo")
        s.add(new)
        await s.commit()
        return new.id


async def test_a_rollback_puts_back_edits_and_deletions_and_removes_new_records(seeded):
    ids, payload, _ = seeded
    new_id = await _change_everything(ids)

    result = await backup.replace_from_backup(payload)

    assert result["mode"] == "rollback"
    async with async_session() as s:
        kept = await s.get(Note, ids["kept"])
        assert (kept.title, kept.body) == ("before", "as backed up")
        assert await s.get(Note, ids["doomed"]) is not None
        assert await s.get(Note, new_id) is None
        titles = (await s.execute(
            select(Project.title).where(Project.user_id == ids["user"])
        )).scalars().all()
        assert titles == ["Rollback drill"]

    after = await backup.export_full_backup()
    for table in ("projects", "milestones", "notes"):
        assert _own(after, table, ids["user"]) == _own(payload, table, ids["user"]), table


async def test_a_surviving_users_api_key_still_works_and_their_sessions_end(seeded):
    ids, payload, key = seeded
    async with async_session() as s:
        version_before = (await s.get(User, ids["user"])).session_version
    await _change_everything(ids)

    await backup.replace_from_backup(payload)

    found = await lookup_key(key)
    assert found is not None and found.user_id == ids["user"]
    async with async_session() as s:
        assert (await s.get(User, ids["user"])).session_version == version_before + 1


async def test_new_records_after_a_rollback_take_fresh_ids(seeded):
    ids, payload, _ = seeded
    await _change_everything(ids)
    await backup.replace_from_backup(payload)

    async with async_session() as s:
        highest = (await s.execute(select(func.max(Note.id)))).scalar_one()
        note = Note(user_id=ids["user"], project_id=ids["project"], title="after rollback",
                    body="", note_type="note", task_kind="work", status="todo")
        s.add(note)
        await s.commit()
        assert note.id > highest
        await s.delete(note)
        await s.commit()


@pytest.mark.parametrize("instance_id", ["another-instance", None])
async def test_a_backup_from_elsewhere_is_refused_on_a_live_instance(seeded, instance_id):
    ids, payload, _ = seeded
    foreign = {**payload, "instance_id": instance_id}
    new_id = await _change_everything(ids)

    with pytest.raises(backup.RestoreRefused):
        await backup.replace_from_backup(foreign)

    async with async_session() as s:
        assert (await s.get(Note, ids["kept"])).title == "after"
        assert await s.get(Note, new_id) is not None
    await backup.replace_from_backup(payload)


async def test_a_users_own_export_is_refused(seeded):
    _, payload, _ = seeded
    with pytest.raises(backup.RestoreRefused):
        await backup.replace_from_backup({**payload, "scope": "user"})


async def test_a_failure_part_way_changes_nothing(seeded):
    ids, payload, _ = seeded
    new_id = await _change_everything(ids)

    with patch.object(backup, "_reset_sequences", AsyncMock(side_effect=RuntimeError("boom"))):
        with pytest.raises(RuntimeError):
            await backup.replace_from_backup(payload)

    async with async_session() as s:
        assert (await s.get(Note, ids["kept"])).title == "after"
        assert await s.get(Note, ids["doomed"]) is None
        assert await s.get(Note, new_id) is not None
    await backup.replace_from_backup(payload)


async def test_a_fresh_instance_adopts_the_backups_instance_id(seeded):
    _, payload, _ = seeded
    own = await get_instance_id()
    try:
        with patch.object(backup, "_is_fresh", AsyncMock(return_value=True)):
            result = await backup.replace_from_backup({**payload, "instance_id": "a-lineage"})
        assert result["mode"] == "fresh"
        assert await get_instance_id() == "a-lineage"
    finally:
        async with async_session() as s:
            await adopt_instance_id(s, own)
            await s.commit()


async def test_the_command_only_previews_without_yes(seeded, tmp_path):
    ids, payload, _ = seeded
    path = write_archive(payload, tmp_path, datetime.now(timezone.utc), app_version="test")
    new_id = await _change_everything(ids)

    assert await run(path) == 4

    async with async_session() as s:
        assert (await s.get(Note, ids["kept"])).title == "after"
        assert await s.get(Note, new_id) is not None
    await backup.replace_from_backup(payload)


async def test_the_command_rolls_back_from_the_tar_with_yes(seeded, tmp_path):
    """The scheduled backup's own file, read back through the command: the
    whole store returns to it, after the safety backup is taken."""
    ids, payload, _ = seeded
    path = write_archive(payload, tmp_path, datetime.now(timezone.utc), app_version="test")
    new_id = await _change_everything(ids)
    saved = {"file": "grimoire-safety.tar", "size_bytes": 1, "where": str(tmp_path)}

    with patch.object(backup_job, "save_state_before_restore",
                      AsyncMock(return_value=saved)) as save:
        assert await run(path, yes=True) == 0

    save.assert_awaited_once()
    async with async_session() as s:
        kept = await s.get(Note, ids["kept"])
        assert (kept.title, kept.body) == ("before", "as backed up")
        assert await s.get(Note, ids["doomed"]) is not None
        assert await s.get(Note, new_id) is None