zhifu-platform/tests/rabbit_cluster_acceptance.py

import asyncio
import json
import secrets
from datetime import UTC, datetime
from pathlib import Path
from time import monotonic
from uuid import UUID, uuid4

import httpx
import pytest
from sqlalchemy import func, select
from test_after_sales import scenario  # noqa: F401
from test_fulfillment import approve, connected, operation_for  # noqa: F401

from zhifu.modules.after_sales.models import BusinessEvent
from zhifu.modules.integrations.fulfillment import execute_due
from zhifu.modules.integrations.models import ConsumerReceipt, ExternalOperation, Reservation
from zhifu.modules.integrations.worker import Broker, consume_one, publish_one
from zhifu.oms.models import Exchange

pytestmark = pytest.mark.integration


class LocalCluster:
    def __init__(self):
        self.identifier = "zhifu-quorum-" + uuid4().hex[:10]
        self.names = [self.identifier + "-" + str(index) for index in range(3)]
        self.password = secrets.token_hex(24)
        self.cookie = secrets.token_hex(24)
        self.report = {
            "execution_id": self.identifier, "started_at": datetime.now(UTC).isoformat(),
            "model_calls": 0, "message_limit": 20, "duration_limit_seconds": 300,
            "memory_bytes_per_node": 512 * 1024 * 1024, "cpu_limit_per_node": 1,
            "published_messages": 0, "consumed_messages": 0, "management_requests": 0,
            "docker_commands": 0, "status": "RUNNING", "phases": [],
        }
        self.started = monotonic()

    async def docker(self, *arguments):
        self.report["docker_commands"] += 1
        process = await asyncio.create_subprocess_exec(
            "docker", *arguments, stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE)
        try:
            async with asyncio.timeout(30):
                output, error = await process.communicate()
        finally:
            if process.returncode is None:
                process.kill()
                await process.wait()
        if process.returncode:
            detail = error.decode("utf-8", errors="replace")
            detail = detail.replace(self.password, "[redacted]").replace(self.cookie, "[redacted]")
            raise RuntimeError(f"Docker {arguments[0]} failed: {detail}")
        return output.decode("utf-8").strip()

    def phase(self, name):
        self.report["phases"].append({"name": name, "elapsed_seconds": round(
            monotonic() - self.started, 3)})

    async def healthy(self, names):
        async with asyncio.timeout(65):
            while True:
                rows = json.loads(await self.docker("inspect", *names))
                if all(row["State"].get("Health", {}).get("Status") == "healthy" for row in rows):
                    return
                if any(row["State"]["Status"] == "exited" for row in rows):
                    raise RuntimeError("临时 RabbitMQ 节点提前退出")
                await asyncio.sleep(2)

    async def start(self):
        await self.docker("network", "create", "--label", "zhifu.exercise=" + self.identifier,
                          self.identifier)
        for index, name in enumerate(self.names):
            # 所有节点和匿名数据卷仅属于本次演练，使用独立网络、端口及受限资源。
            await self.docker(
                "run", "-d", "--name", name, "--hostname", "rq" + str(index),
                "--network", self.identifier, "--label", "zhifu.exercise=" + self.identifier,
                "--memory", "512m", "--memory-swap", "512m", "--cpus", "1",
                "--pids-limit", "160", "-p", "127.0.0.1::5672", "-p", "127.0.0.1::15672",
                "-e", "RABBITMQ_DEFAULT_USER=exercise",
                "-e", "RABBITMQ_DEFAULT_PASS=" + self.password,
                "-e", "RABBITMQ_ERLANG_COOKIE=" + self.cookie,
                "-e", "RABBITMQ_SERVER_ADDITIONAL_ERL_ARGS=+S 2:2 +SDcpu 1 +SDio 1",
                "-e", "RABBITMQ_CTL_ERL_ARGS=+S 1:1 +SDcpu 1 +SDio 1",
                "--health-cmd", "rabbitmq-diagnostics -q check_running",
                "--health-interval", "2s", "--health-timeout", "5s", "--health-retries", "30",
                "rabbitmq:4.2-management",
            )
        await self.healthy(self.names)
        for name in self.names[1:]:
            await self.docker("exec", name, "rabbitmqctl", "join_cluster", "rabbit@rq0")
        await self.healthy(self.names)
        rows = await self.refresh_ports()
        self.report["image_id"] = rows[0]["Image"]
        self.phase("three_nodes_joined")

    async def refresh_ports(self):
        # Docker 的随机发布端口在容器再次启动后重新读取，不缓存旧的主机端口。
        rows = json.loads(await self.docker("inspect", *self.names))
        self.ports = {
            row["Name"].removeprefix("/"): {
                key: int(value[0]["HostPort"])
                for key, value in row["NetworkSettings"]["Ports"].items() if value
            } for row in rows
        }
        return rows

    def amqp(self, name):
        return f"amqp://exercise:{self.password}@127.0.0.1:{self.ports[name]['5672/tcp']}/"

    async def query(self, name, path):
        self.report["management_requests"] += 1
        port = self.ports[name]["15672/tcp"]
        async with httpx.AsyncClient(timeout=8, auth=("exercise", self.password)) as client:
            response = await client.get(f"http://127.0.0.1:{port}/api/{path}")
            response.raise_for_status()
            return response.json()

    async def cleanup(self):
        # 只移除带有本次随机标签的资源，匿名卷随所属容器一起移除。
        container_ids = (await self.docker(
            "ps", "-aq", "--filter", "label=zhifu.exercise=" + self.identifier)).splitlines()
        volumes = set()
        if container_ids:
            rows = json.loads(await self.docker("inspect", *container_ids))
            volumes = {mount["Name"] for row in rows for mount in row["Mounts"]
                       if mount["Type"] == "volume"}
            await self.docker("rm", "-fv", *container_ids)
        network_ids = (await self.docker(
            "network", "ls", "-q", "--filter", "label=zhifu.exercise=" + self.identifier
        )).splitlines()
        if network_ids:
            await self.docker("network", "rm", *network_ids)
        remaining = await self.docker(
            "ps", "-aq", "--filter", "label=zhifu.exercise=" + self.identifier)
        networks = await self.docker(
            "network", "ls", "-q", "--filter", "label=zhifu.exercise=" + self.identifier)
        present_volumes = set((await self.docker("volume", "ls", "-q")).splitlines())
        assert not remaining and not networks and not (volumes & present_volumes)
        self.report["removed_anonymous_volumes"] = len(volumes)
        self.report["cleanup_verified"] = True
        self.phase("temporary_resources_removed")


async def test_three_node_quorum_leader_failure_and_business_uniqueness(connected):  # noqa: F811
    cluster = LocalCluster()
    report = cluster.report
    output = Path(".work/rabbit-cluster-20260922.json")
    try:
        # 业务演练最多 235 秒，另外保留 55 秒处理所有临时资源清理。
        async with asyncio.timeout(235):
            await cluster.start()
            nodes = await cluster.query(cluster.names[0], "nodes")
            assert len(nodes) == 3 and all(node["running"] for node in nodes)
            report["nodes_before"] = [node["name"] for node in nodes]
            ticket, event_id = await approve(connected)
            sessions = connected["app"].state.sessions
            queue = "exercise." + uuid4().hex
            report["queue"] = queue
            report["event_id"] = str(UUID(bytes=event_id))
            report["ticket_id"] = ticket["ticket_id"]
            async with Broker(cluster.amqp(cluster.names[0]), queue, queue) as broker:
                assert await publish_one(sessions, broker, event_id)
                report["published_messages"] += 1
                async with sessions() as session:
                    event = await session.get(BusinessEvent, event_id)
                    assert event.delivery_status == "SENT"
                    await broker.publish_event(event)
                    report["published_messages"] += 1
            before = await cluster.query(cluster.names[0], "queues/%2F/" + queue)
            assert before["type"] == "quorum" and len(before["members"]) == 3
            leader = before["leader"]
            report["queue_before"] = {key: before[key] for key in ("type", "leader", "members")}
            # 根据实际 leader 确定被终止节点，确保验证包含选举。
            index = int(leader.removeprefix("rabbit@rq"))
            stopped = cluster.names[index]
            survivor = next(name for name in cluster.names if name != stopped)
            cluster.phase("confirmed_messages_before_failure")
            await cluster.docker("kill", stopped)
            failure_started = monotonic()
            cluster.phase("leader_terminated")
            async with asyncio.timeout(30):
                while True:
                    current = await cluster.query(survivor, "queues/%2F/" + queue)
                    if current.get("leader") and current["leader"] != leader:
                        break
                    await asyncio.sleep(1)
            report["queue_after_failure"] = {key: current[key] for key in (
                "type", "leader", "members")}
            async with Broker(cluster.amqp(survivor), queue, queue) as broker:
                for _ in range(2):
                    assert await consume_one(sessions, broker, 5)
                    report["consumed_messages"] += 1
                assert await consume_one(sessions, broker, 5) is False
            report["recovery_seconds"] = round(monotonic() - failure_started, 3)
            cluster.phase("survivor_consumed_confirmed_messages")
            operation = await operation_for(connected, ticket)
            assert operation is not None
            assert await execute_due(sessions, 5, operation.id)
            async with sessions() as session:
                for model in (ExternalOperation, Reservation, ConsumerReceipt):
                    count = await session.scalar(select(func.count()).select_from(model).where(
                        model.tenant_id == connected["tenant"].bytes))
                    assert count == 1
            async with connected["oms_sessions"]() as session:
                assert await session.scalar(select(func.count()).select_from(Exchange).where(
                    Exchange.tenant_id == connected["tenant"].bytes)) == 1
            report["external_operations"] = report["oms_exchanges"] = 1
            await cluster.docker("start", stopped)
            await cluster.healthy([stopped])
            await cluster.refresh_ports()
            nodes = await cluster.query(survivor, "nodes")
            assert len(nodes) == 3 and all(node["running"] for node in nodes)
            report["nodes_recovered"] = [node["name"] for node in nodes]
            async with Broker(cluster.amqp(stopped), queue, queue) as broker:
                async with sessions() as session:
                    event = await session.get(BusinessEvent, event_id)
                    await broker.publish_event(event)
                    report["published_messages"] += 1
                assert await consume_one(sessions, broker, 5)
                report["consumed_messages"] += 1
            async with sessions() as session:
                count = await session.scalar(select(func.count()).select_from(
                    ExternalOperation).where(
                        ExternalOperation.tenant_id == connected["tenant"].bytes))
                assert count == 1
            cluster.phase("original_node_recovered_and_duplicate_consumed")
            report["status"] = "PASSED"
    except Exception as failure:
        report["status"] = "FAILED"
        report["failure"] = f"{type(failure).__name__}: {failure}"
        raise
    finally:
        try:
            async with asyncio.timeout(55):
                await cluster.cleanup()
        except Exception as failure:
            report["status"] = "FAILED"
            report["cleanup_verified"] = False
            report["cleanup_failure"] = f"{type(failure).__name__}: {failure}"
            raise
        finally:
            report["elapsed_seconds"] = round(monotonic() - cluster.started, 3)
            await asyncio.to_thread(output.parent.mkdir, parents=True, exist_ok=True)
            await asyncio.to_thread(output.write_text, json.dumps(
                report, ensure_ascii=False, indent=2), encoding="utf-8")