site/apps/api/tests/orchestrator_dual_harness.py

"""Orchestrator + two virtual hub clients for dual-hub sync E2E tests."""

from __future__ import annotations

import copy
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import Any
from uuid import uuid4

from fastapi.testclient import TestClient

from app.modules.sync import repository as repo


@dataclass
class HubClient:
    hub_site_id: str
    farm_hub_id: str
    api_key: str
    pull_cursor: int = 0
    local_catalog: dict[str, dict[str, dict[str, Any]]] = field(default_factory=dict)

    def auth_header(self) -> dict[str, str]:
        return {"Authorization": f"Hub {self.hub_site_id}:{self.api_key}"}


class OrchestratorDualHubHarness:
    def __init__(self, client: TestClient, enterprise_id: str) -> None:
        self.client = client
        self.enterprise_id = enterprise_id
        self.hub_a: HubClient | None = None
        self.hub_b: HubClient | None = None

    def pair_hubs(self) -> tuple[HubClient, HubClient]:
        suffix = uuid4().hex[:8]
        code_a = self._start_pairing()
        hub_a = self._confirm_hub(code_a, f"hub-site-a-{suffix}", "Hub A")
        code_b = self._start_pairing()
        hub_b = self._confirm_hub(code_b, f"hub-site-b-{suffix}", "Hub B")
        self.hub_a = hub_a
        self.hub_b = hub_b
        return hub_a, hub_b

    def _start_pairing(self) -> str:
        login = self.client.post(
            "/api/v1/auth/login",
            json={"email": "admin@compton.example", "password": "Admin1234"},
        )
        token = login.json()["access_token"]
        start = self.client.post(
            "/api/v1/enterprise/pair/start",
            headers={"Authorization": f"Bearer {token}"},
            json={"enterprise_id": self.enterprise_id, "farm_name": "Farm"},
        )
        return start.json()["code"]

    def _confirm_hub(self, code: str, hub_site_id: str, hub_name: str) -> HubClient:
        confirm = self.client.post(
            "/api/v1/enterprise/pair/confirm",
            json={"code": code, "hub_site_id": hub_site_id, "hub_name": hub_name},
        )
        body = confirm.json()
        return HubClient(
            hub_site_id=body["hub_site_id"],
            farm_hub_id=body["farm_hub_id"],
            api_key=body["api_key"],
        )

    def edit_local(self, hub: HubClient, table: str, record_id: str, payload: dict, *, version: int, content_hash: str) -> dict:
        row = copy.deepcopy(payload)
        row["id"] = record_id
        row["version"] = version
        row["content_hash"] = content_hash
        hub.local_catalog.setdefault(table, {})[record_id] = row
        return row

    def push_from(self, hub: HubClient, events: list[dict]) -> dict:
        resp = self.client.post(
            "/api/v1/sync/changes/push",
            headers=hub.auth_header(),
            json={"events": events},
        )
        resp.raise_for_status()
        return resp.json()

    def pull_for(self, hub: HubClient) -> list[dict]:
        resp = self.client.post(
            "/api/v1/sync/changes/pull",
            headers=hub.auth_header(),
            json={"cursor": hub.pull_cursor, "limit": 100},
        )
        resp.raise_for_status()
        body = resp.json()
        events = body.get("events") or []
        ack_ids: list[str] = []
        for ev in events:
            self._apply_local(hub, ev)
            ack_ids.append(ev["event_id"])
        if ack_ids:
            self.client.post(
                "/api/v1/sync/changes/ack",
                headers=hub.auth_header(),
                json={"event_ids": ack_ids, "direction": "inbound"},
            )
        if body.get("next_cursor") is not None:
            hub.pull_cursor = int(body["next_cursor"])
        return events

    def _apply_local(self, hub: HubClient, event: dict) -> None:
        table = event["table"]
        record_id = event["record_id"]
        hub.local_catalog.setdefault(table, {})[record_id] = {
            **(event.get("payload") or {}),
            "id": record_id,
            "version": event.get("version"),
            "content_hash": event.get("content_hash"),
        }

    def drain(self, hubs: list[HubClient] | None = None, rounds: int = 10) -> None:
        targets = hubs or [h for h in (self.hub_a, self.hub_b) if h]
        for _ in range(rounds):
            for hub in targets:
                self.pull_for(hub)

    def get_catalog(self, hub: HubClient, table: str, record_id: str) -> dict | None:
        return hub.local_catalog.get(table, {}).get(record_id)

    def get_orchestrator_catalog(self, table: str, record_id: str) -> dict | None:
        from app.modules.zootech.catalog_apply import load_catalog_row

        return load_catalog_row(self.enterprise_id, table, record_id)

    def make_component_event(
        self,
        hub: HubClient,
        record_id: str,
        *,
        name: str,
        version: int,
        content_hash: str,
        event_id: str | None = None,
        dry_matter: float = 88.0,
    ) -> dict:
        payload = {
            "name": name,
            "type": "grain",
            "dry_matter": dry_matter,
            "protein": 8.0,
            "energy": 1.2,
            "price": 0.0,
            "is_active": True,
        }
        self.edit_local(hub, "component", record_id, payload, version=version, content_hash=content_hash)
        return {
            "event_id": event_id or str(uuid4()),
            "seq": version,
            "domain": "global",
            "table": "component",
            "record_id": record_id,
            "action": "upsert",
            "version": version,
            "content_hash": content_hash,
            "payload": {**payload, "id": record_id, "version": version, "content_hash": content_hash},
            "emitted_at": datetime.now(UTC).isoformat(),
            "origin_site_id": hub.hub_site_id,
        }


def make_enterprise(client: TestClient) -> str:
    slug = f"test-ent-{uuid4().hex[:8]}"
    ent = repo.create_enterprise("Test Enterprise", slug)
    from app.modules.users.repository import get_user_by_email

    user = get_user_by_email("admin@compton.example")
    assert user
    repo.add_member(user.id, ent.id, "admin")
    return ent.id