zhifu-platform/tests/redis_recovery_acceptance.py

import asyncio
import importlib.util
import json
import time
from pathlib import Path
from uuid import uuid4

import pytest
from sqlalchemy import select
from test_after_sales import scenario  # noqa: F401

from zhifu.modules.agent_runtime.capacity import acquire, keys, release
from zhifu.modules.agent_runtime.config import AgentSettings
from zhifu.modules.memory.models import MemoryProfile


async def test_actual_memory_snapshot_and_isolated_recovery(scenario):  # noqa: F811
    spec = importlib.util.spec_from_file_location("redis_recovery", "scripts/redis_recovery.py")
    recovery = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(recovery)
    call, state = scenario["call"], scenario["app"].state.state
    path = "/openapi/v1/me/memory"
    response = await call(
        "PUT", path + "/consent", payload={"expected_revision": 0, "enabled": True}
    )
    assert response.status_code == 200
    response = await call(
        "PUT",
        path,
        payload={
            "expected_revision": 1,
            "communication": {"response_style": "concise", "language": "zh-CN"},
        },
    )
    assert response.status_code == 200, response.text
    async with scenario["app"].state.sessions() as session:
        row = await session.scalar(
            select(MemoryProfile).where(
                MemoryProfile.tenant_id == scenario["tenant"].bytes,
                MemoryProfile.subject_id == scenario["accounts"]["customer"].bytes,
            )
        )
        value_key = row.value_key
    found = []
    async for key in state.scan_iter(match="zhifu_memory:*"):
        content = await state.json().get(key)
        if content["key"] == value_key:
            found.append((key.decode(), content))
    assert len(found) == 1
    key, expected = found[0]
    expiry = await state.pexpiretime(key)
    assert expiry > int(time.time() * 1000)
    transient = "monitor:worker:recovery:" + uuid4().hex
    expiring = "recovery:expires:" + uuid4().hex
    await state.set(transient, "active", ex=120)
    await state.set(expiring, "expires", ex=15)
    expires_at = await state.pexpiretime(expiring)
    capacity_id = uuid4().hex
    assert await acquire(state, AgentSettings(), "background", capacity_id)
    interactive_id = uuid4().hex
    assert await acquire(
        state, AgentSettings(), "interactive", interactive_id, tenant=scenario["tenant"].hex
    )
    try:
        directory = await asyncio.to_thread(recovery.backup)
        assert time.time() * 1000 < expires_at, "快照完成前短期记录已过期，需重新安排验收时间"
        # RDB 保留绝对到期时间，恢复过程不能重新开始原 TTL。
        await asyncio.sleep(max(0, expires_at / 1000 - time.time()) + 0.2)

        def inspect_recovery():
            with recovery.recovered(directory) as (container, entry):
                assert json.loads(recovery.cli(container, "JSON.GET", key)) == expected
                assert recovery.cli(container, "PEXPIRETIME", key) == expiry
                assert recovery.cli(container, "EXISTS", expiring) == 0
                assert recovery.cli(container, "GET", transient) == "active"
                assert (
                    recovery.cli(container, "ZSCORE", "agent:active:global:background", capacity_id)
                    is not None
                )
                assert (
                    recovery.cli(container, "HGET", "agent:active:global:owners", interactive_id)
                    == scenario["tenant"].hex
                )
                query = "@key:{" + value_key + "}"
                result = recovery.cli(container, "FT.SEARCH", "zhifu_memory", query, "NOCONTENT")
                assert result["total_results"] == 1, result
                removed = recovery.clear_transient(container)
                assert recovery.cli(container, "EXISTS", transient) == 0
                assert recovery.cli(container, "EXISTS", *keys("agent:active:global")) == 0
                assert json.loads(recovery.cli(container, "JSON.GET", key)) == expected
                return {"indexes": entry["indexes"], "transient_removed": removed}

        report = await asyncio.to_thread(inspect_recovery)
        source = await asyncio.to_thread(
            recovery.docker,
            "compose",
            "ps",
            "-q",
            "redis-state",
            cwd=recovery.ROOT,
            capture_output=True,
        )
        with pytest.raises(ValueError, match="独立恢复容器"):
            await asyncio.to_thread(recovery.clear_transient, source.stdout.decode().strip())
        assert await state.get(transient) == b"active"
        assert (await call("GET", path)).json()["content"] == response.json()["content"]
        backup_path = directory / "state.rdb"
        original = await asyncio.to_thread(backup_path.read_bytes)
        try:
            await asyncio.to_thread(backup_path.write_bytes, original + b"invalid")
            with pytest.raises(ValueError, match="SHA256"):
                await asyncio.to_thread(recovery.verify, directory)
        finally:
            await asyncio.to_thread(backup_path.write_bytes, original)
        report.update(
            memory_equal=True,
            absolute_expiry_equal=True,
            expired_key_absent=True,
            search_index_query=True,
            source_unchanged=True,
            corrupted_backup_rejected=True,
            source_cleanup_rejected=True,
            model_pool_leases_cleared=True,
        )
        await asyncio.to_thread(
            Path("docs/verification/redis-recovery.json").write_text,
            json.dumps(report, indent=2),
            encoding="utf-8",
        )
    finally:
        await release(state, capacity_id)
        await release(state, interactive_id)
        await state.delete(transient, expiring)