chuanqi/cq2_sdk/scripts/demo_server_zone_store.py

"""End-to-end demo: B1 (find_server) + B2 (bulk_upsert zones) + B3 (pick latest).

This script simulates the 996box daily zone refresh and demonstrates the
B1 / B2 / B3 business flow end-to-end against the real MySQL.

Flow:
  B1: ServerStore.upsert("玖扬复古176", 14414738)
      → returns Server with id
  B2: ZoneStore.bulk_upsert(server_id, [9 zones for today])
      → writes 9 rows
      → ZoneStore.mark_missing_closed(server_id, kept_ids)
        (simulates zones that 996box no longer reports → close them)
  B3: ZoneStore.get_latest_openable(server_id)
      → returns the most recent zone that is status=open AND open_time <= NOW()

Run:
  uv run python scripts/demo_server_zone_store.py
  uv run python scripts/demo_server_zone_store.py --cleanup
"""
from __future__ import annotations
import argparse
import sys
from datetime import datetime, timedelta
from pathlib import Path

_HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(_HERE.parent / "src"))

from cq2_sdk.env_loader import load_env  # noqa: E402
_PROJECT_ROOT = _HERE.parent
load_env(_PROJECT_ROOT / ".env")

from cq2_sdk.db_models import Zone as DbZone  # noqa: E402
from cq2_sdk.store import ServerStore, ZoneStore, connection  # noqa: E402

TEST_SERVER = "DEMO_玖扬复古176"
TEST_SERVER_ID_996BOX = 14414738
GAME_ID = 206923

# 9 daily time slots for 玖扬复古176 (from business-tasks.md §4 B2)
TIME_SLOTS = [
    (0, 0), (7, 0), (10, 0), (12, 0), (14, 0),
    (16, 0), (19, 0), (20, 30), (21, 30),
]


def build_todays_zones() -> list[DbZone]:
    """Construct Zone objects for today's 9 standard time slots.

    All marked 'open' and given open_time earlier than NOW so B3 can pick one.
    """
    now = datetime.now().replace(microsecond=0)
    zones: list[DbZone] = []
    for i, (h, m) in enumerate(TIME_SLOTS):
        open_time = now.replace(hour=h, minute=m, second=0)
        # If today's slot is in the future, use a slot from earlier today
        if open_time > now:
            open_time = now - timedelta(hours=i + 1)
        zones.append(DbZone(
            zone_id=f"slot_{h:02d}{m:02d}",
            zone_name=f"玖扬复古176 / {h:02d}:{m:02d}",
            host="jycq-206923.dhsf.xqhuyu.com",
            port=8001,
            open_time=open_time,
            status="open",
        ))
    return zones


def main() -> int:
    parser = argparse.ArgumentParser(description="Demo B1/B2/B3 server+zone flow")
    parser.add_argument("--cleanup", action="store_true",
                        help="Delete the test server and its zones after demo")
    args = parser.parse_args()

    print("=" * 70)
    print("B1/B2/B3 end-to-end demo (ServerStore + ZoneStore)")
    print("=" * 70)

    with connection() as conn:
        servers = ServerStore(conn)
        zones = ZoneStore(conn)

        # ---------- B1: find / upsert server ----------
        print("\n[B1] ServerStore.upsert(玖扬复古176, 14414738)")
        server = servers.upsert(TEST_SERVER, TEST_SERVER_ID_996BOX, game_id=GAME_ID)
        print(f"      server.id        = {server.id}")
        print(f"      server.server_id = {server.server_id} (996box id)")
        print(f"      server.first_seen_at = {server.first_seen_at}")

        # Read-back via both lookup methods
        refetched = servers.get_by_name(TEST_SERVER)
        assert refetched is not None and refetched.id == server.id
        by_996box = servers.get_by_server_id(TEST_SERVER_ID_996BOX)
        assert by_996box is not None and by_996box.id == server.id
        print(f"      get_by_name     -> id={refetched.id}  [OK]")
        print(f"      get_by_server_id({TEST_SERVER_ID_996BOX}) -> id={by_996box.id}  [OK]")

        # ---------- B2: bulk_upsert today's zones ----------
        print(f"\n[B2] ZoneStore.bulk_upsert(server.id={server.id}, 9 zones)")
        todays_zones = build_todays_zones()
        n = zones.bulk_upsert(server.id, todays_zones)
        print(f"      rows affected: {n}")

        # Simulate 996box removing one zone: drop "slot_0700" and refresh
        kept_ids = {z.zone_id for z in todays_zones if z.zone_id != "slot_0700"}
        closed = zones.mark_missing_closed(server.id, kept_ids)
        print(f"      mark_missing_closed (simulated 996box dropped 'slot_0700'): "
              f"{closed} zone(s) flipped to 'closed'")

        # Verify today's open list
        todays = zones.get_today_open(server.id)
        print(f"      get_today_open → {len(todays)} zone(s) for today")
        for z in todays[:5]:
            print(f"        - {z.zone_id}  open={z.open_time}  status={z.status}")
        if len(todays) > 5:
            print(f"        ... ({len(todays) - 5} more)")

        # ---------- B3: pick the latest openable zone ----------
        print("\n[B3] ZoneStore.get_latest_openable(server.id)")
        latest = zones.get_latest_openable(server.id)
        if latest is None:
            print("      no openable zone found (B3 would raise NoZoneAvailable)")
        else:
            print(f"      latest.zone_id    = {latest.zone_id}")
            print(f"      latest.open_time  = {latest.open_time}")
            print(f"      latest.status     = {latest.status}")
            print(f"      latest.host:port  = {latest.host}:{latest.port}")

        # ---------- summary ----------
        print("\n[summary] persisted state in MySQL 996cq:")
        print(f"    servers.name = '{server.server_name}'  id={server.id}")
        print(f"    zones       = {len(todays)} today, {n} total upserted")

        if args.cleanup:
            print("\n[cleanup] deleting test server and zones...")
            with conn.cursor() as cur:
                cur.execute("DELETE FROM zones WHERE server_id = %s", (server.id,))
                cur.execute("DELETE FROM servers WHERE id = %s", (server.id,))
            print("    done")
        else:
            print("\n  (use --cleanup to delete the test rows)")

    print("\n[OK] demo complete.")
    return 0


if __name__ == "__main__":
    sys.exit(main())