diff --git a/yunke-server/scripts/__init__.py b/yunke-server/scripts/__init__.py new file mode 100644 index 0000000..3af55df --- /dev/null +++ b/yunke-server/scripts/__init__.py @@ -0,0 +1,7 @@ +"""Operational scripts for the yunke data mirror. + +Each module under this package is a standalone, runnable integration check or +maintenance helper (e.g. the Task 1.9 smoke script). Importing this package has +no side effects — the scripts are entry-point functions meant to be invoked via +``python -m scripts.`` from the ``yunke-server/`` directory. +""" diff --git a/yunke-server/scripts/smoke_wechat.py b/yunke-server/scripts/smoke_wechat.py new file mode 100644 index 0000000..974b92b --- /dev/null +++ b/yunke-server/scripts/smoke_wechat.py @@ -0,0 +1,358 @@ +"""End-to-end smoke check for the yunke WeChat data mirror (Task 1.9). + +Wires the existing pieces together against the REAL yunke CRM (read-only) and a +local SQLite DB, then prints a summary. It exists to prove the plumbing +(session -> scraper -> upsert -> DB) holds end-to-end before the full sync +pipeline is built. + +Auth modes (mutually exclusive; session mode preferred for now) +-------------------------------------------------------------- +1. **Existing-session mode (preferred today).** Set ``YK_USERID``, + ``YK_COMPANY``, and optionally ``YK_COOKIE``. The script builds a + :class:`YunkeClient` directly from that captured session + (``userid=`` / ``company=`` / ``cookies=``) and BYPASSES login. This is the + only auth mode that actually authenticates against the real CRM today — + see the Phase-0 note below. + +2. **Login mode (future).** Set ``YUNKE_USERNAME`` and ``YUNKE_PASSWORD`` (and + leave the session vars unset). The script calls :meth:`YunkeClient.login`, + which runs the ``getSecureKey`` -> ``pcLogin`` flow. **NOTE:** password + encryption is a Phase-0 placeholder (see + :func:`app.session.login.encrypt_password`), so the real CRM will REJECT the + login. The script catches that, prints a clear message, and exits with code + 3. Once Phase-0 login lands, this mode will just work. + + If NEITHER set of env vars is present, the script prints usage and exits + with code 2. + +What it does (kept tiny — read-only) +------------------------------------ +* Creates all ORM tables on the smoke DB (``Base.metadata.create_all``). +* Scrapes up to ~10 WeChat friends (:func:`scrape_wechat_friends`, first page + only), upserts them via :func:`upsert_customers`, prints the count and the + first friend (friend_id / friend_nickname / owner_account_id). +* For the first scraped friend, scrapes up to ~50 of their messages from the + last 7 days (:func:`scrape_wechat_messages`), upserts via + :func:`upsert_messages`, prints the count and a 3-row sample. +* Verifies rows landed by querying the smoke DB back, and prints a final + summary. + +The script NEVER writes to the upstream CRM — every upstream call is a list / +search endpoint. It writes only to the local smoke DB. + +How to run +---------- +From ``yunke-server/`` with the venv active (``.venv/Scripts/activate`` on +Windows, ``source .venv/bin/activate`` on *nix):: + + # Existing-session mode (works today — controller supplies the values): + export YK_USERID=aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + export YK_COMPANY=8m0atd + export YK_COOKIE="JSESSIONID=deadbeef; other=foo" + python -m scripts.smoke_wechat + + # Login mode (Phase-0 blocked — will exit 3 with a message): + export YUNKE_USERNAME=user + export YUNKE_PASSWORD=pass + python -m scripts.smoke_wechat + +The smoke DB defaults to ``./smoke.db`` (SQLite); override the URL with +``SMOKE_DB_URL`` (e.g. ``sqlite+aiosqlite:///./smoke.db`` or any other +SQLAlchemy async URL). On a 401 / expired session, the script prints a helpful +message and exits non-zero (code 4) instead of dumping a traceback. + +Exit codes: ``0`` success · ``2`` no auth env · ``3`` login failed (Phase-0 +blocked) · ``4`` upstream auth/session error mid-scrape. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys +import time +from typing import Any + +import httpx +import sqlalchemy as sa +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +from app.config import get_settings +from app.db import models # noqa: F401 (register tables on Base.metadata) +from app.db.base import Base +from app.db.models import Customer, Message +from app.scraper.wechat import scrape_wechat_friends, scrape_wechat_messages +from app.session.yunke_client import YunkeClient +from app.storage.repos import upsert_customers, upsert_messages + +# Caps that keep the smoke run tiny (≈10 customers, ≈50 messages). +MAX_FRIENDS = 10 +MAX_MESSAGES = 50 +# Last 7 days in ms — backward window for the message scraper. +MESSAGE_WINDOW_MS = 7 * 86_400_000 + +DEFAULT_SMOKE_DB_URL = "sqlite+aiosqlite:///./smoke.db" + + +def _parse_cookie_env(raw: str) -> dict[str, str]: + """Parse the ``YK_COOKIE`` env var into a ``{name: value}`` dict. + + Accepts either: + + * a raw ``Cookie:`` header (``"JSESSIONID=abc; foo=bar"``), or + * a JSON object (``'{"JSESSIONID": "abc"}'``). + + Empty / unparsable input yields ``{}`` (the caller is fine with no cookies + — the ``userid`` + ``company`` headers may be enough on their own). + """ + raw = (raw or "").strip() + if not raw: + return {} + # JSON first — handy for structured config. + if raw.startswith("{"): + try: + obj = json.loads(raw) + except ValueError: + return {} + if not isinstance(obj, dict): + return {} + return {str(k): str(v) for k, v in obj.items()} + # Otherwise treat as a raw header (semicolon-separated name=value pairs). + out: dict[str, str] = {} + for part in raw.split(";"): + part = part.strip() + if not part or "=" not in part: + continue + name, _, value = part.partition("=") + out[name.strip()] = value.strip() + return out + + +def build_client_from_env() -> YunkeClient | None: + """Construct a :class:`YunkeClient` from env vars (pure / network-free). + + Selection rules (in priority order): + + 1. **Existing-session mode** — if EITHER of ``YK_USERID`` / ``YK_COMPANY`` + is set, BOTH must be set (``YK_COOKIE`` optional). Constructs the client + with that session and bypasses login. + 2. **Login mode** — if ``YUNKE_USERNAME`` and ``YUNKE_PASSWORD`` are both + set, constructs a credentialed client. The caller must call + ``await client.login()`` to actually authenticate. + 3. Otherwise returns ``None`` so the caller can print usage and exit. + + Raises ``ValueError`` if only one of ``YK_USERID`` / ``YK_COMPANY`` is set + (a clear caller bug, not something to silently paper over). + + Pure / network-free — safe to unit-test with ``monkeypatch.setenv``. + """ + base_url = os.environ.get("YUNKE_BASE_URL") or get_settings().yunke_base_url + + userid = os.environ.get("YK_USERID", "").strip() + company = os.environ.get("YK_COMPANY", "").strip() + cookie_raw = os.environ.get("YK_COOKIE", "") + + if userid or company: + # Session mode — BOTH required. + if not (userid and company): + raise ValueError( + "YK_USERID and YK_COMPANY must be set together " + "(got only one)." + ) + return YunkeClient( + base_url=base_url, + userid=userid, + company=company, + cookies=_parse_cookie_env(cookie_raw), + ) + + username = os.environ.get("YUNKE_USERNAME", "").strip() + password = os.environ.get("YUNKE_PASSWORD", "") + if username and password: + return YunkeClient( + base_url=base_url, + username=username, + password=password, + ) + + return None + + +def _print_usage(stream=sys.stderr) -> None: + stream.write( + "smoke_wechat: missing auth env vars.\n" + "Set EITHER existing-session env (YK_USERID + YK_COMPANY [+ YK_COOKIE])\n" + " OR login env (YUNKE_USERNAME + YUNKE_PASSWORD).\n" + "See `python -m scripts.smoke_wechat --help` or the module docstring.\n" + ) + + +async def main() -> int: + """Run the smoke check. Returns a process exit code (0 = success).""" + client = build_client_from_env() + if client is None: + _print_usage() + return 2 + + # Login mode: try to log in. Phase-0 placeholder encryption will be + # rejected by the real CRM — catch and explain instead of tracebacking. + if not client.logged_in: + try: + await client.login() + except Exception as exc: # noqa: BLE001 — any failure is fatal-but-explained + sys.stderr.write( + f"\n[smoke] login failed (likely Phase-0 placeholder " + f"encryption): {exc!r}\n" + "Use existing-session mode (YK_USERID / YK_COMPANY / YK_COOKIE) " + "until Phase-0 login lands.\n" + ) + await client.aclose() + return 3 + + smoke_db_url = os.environ.get("SMOKE_DB_URL", DEFAULT_SMOKE_DB_URL) + engine = create_async_engine(smoke_db_url, future=True) + session_maker = async_sessionmaker(engine, expire_on_commit=False) + + try: + # 1) Create all ORM tables on the smoke DB. + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + customers_upserted = 0 + messages_upserted = 0 + first_friend: dict[str, Any] | None = None + sample_msgs: list[dict[str, Any]] = [] + + # 2) Customers — consume only the first page worth (~10 records). + print( + f"[smoke] scraping up to {MAX_FRIENDS} friends from " + f"{client.base_url} ..." + ) + friends: list[dict[str, Any]] = [] + try: + agen = scrape_wechat_friends(client, page_size=MAX_FRIENDS) + try: + async for rec in agen: + friends.append(rec) + if len(friends) >= MAX_FRIENDS: + break + finally: + # Explicit cleanup — `async for` does not synchronously close + # the generator on `break`. + await agen.aclose() + except httpx.HTTPStatusError as exc: + sys.stderr.write( + f"\n[smoke] upstream auth/session error scraping friends: " + f"HTTP {exc.response.status_code}.\n" + "Session likely expired — refresh YK_COOKIE / YK_USERID / " + "YK_COMPANY and retry.\n" + ) + return 4 + + if friends: + async with session_maker() as session: + customers_upserted = await upsert_customers(session, friends) + await session.commit() + first_friend = friends[0] + + if first_friend is None: + print("[smoke] no friends returned; nothing more to do.") + _print_summary(session_maker, customers_upserted, messages_upserted) + return 0 + + print(f"[smoke] customers upserted: {customers_upserted}") + print( + "[smoke] first friend: " + f"friend_id={first_friend.get('friend_id')!r} " + f"friend_nickname={first_friend.get('friend_nickname')!r} " + f"owner_account_id={first_friend.get('owner_account_id')!r}" + ) + + # 3) Messages for the first friend — last 7 days, capped at 50. + owner = first_friend.get("owner_account_id") or "" + friend_id = first_friend.get("friend_id") or "" + if not (owner and friend_id): + print( + "[smoke] first friend missing owner_account_id/friend_id; " + "skipping messages." + ) + else: + end_ms = int(time.time() * 1000) + start_ms = end_ms - MESSAGE_WINDOW_MS + print( + f"[smoke] scraping up to {MAX_MESSAGES} messages for " + f"owner={owner!r} friend={friend_id!r} over the last 7 days ..." + ) + msgs: list[dict[str, Any]] = [] + try: + mgen = scrape_wechat_messages( + client, + wechat_id=owner, + friend_id=friend_id, + start=start_ms, + end=end_ms, + ) + try: + async for rec in mgen: + msgs.append(rec) + if len(msgs) >= MAX_MESSAGES: + break + finally: + await mgen.aclose() + except httpx.HTTPStatusError as exc: + sys.stderr.write( + f"\n[smoke] upstream auth/session error scraping messages: " + f"HTTP {exc.response.status_code}.\n" + "Session likely expired — refresh and retry.\n" + ) + return 4 + + if msgs: + async with session_maker() as session: + messages_upserted = await upsert_messages(session, msgs) + await session.commit() + sample_msgs = msgs[:3] + + print(f"[smoke] messages upserted: {messages_upserted}") + for rec in sample_msgs: + text = (rec.get("text") or "") + print( + f" msg_id={rec.get('msg_id')!r} type={rec.get('type')!r} " + f"type_raw={rec.get('type_raw')!r} ts={rec.get('timestamp')!r} " + f"text={text[:40]!r}" + ) + + # 4) Verify rows landed by querying the smoke DB back. + _print_summary(session_maker, customers_upserted, messages_upserted) + return 0 + finally: + await client.aclose() + await engine.dispose() + + +async def _print_summary( + session_maker: async_sessionmaker, + customers_upserted: int, + messages_upserted: int, +) -> None: + """Print the final summary, including DB row counts read back from SQLite.""" + async with session_maker() as session: + db_customers = await session.scalar( + sa.select(sa.func.count()).select_from(Customer) + ) + db_messages = await session.scalar( + sa.select(sa.func.count()).select_from(Message) + ) + + print("-" * 60) + print("[smoke] SUMMARY") + print(f" customers upserted : {customers_upserted}") + print(f" messages upserted : {messages_upserted}") + print(f" customer rows in DB : {db_customers}") + print(f" message rows in DB : {db_messages}") + print("-" * 60) + + +if __name__ == "__main__": + sys.exit(asyncio.run(main())) diff --git a/yunke-server/tests/test_smoke_script.py b/yunke-server/tests/test_smoke_script.py new file mode 100644 index 0000000..edba803 --- /dev/null +++ b/yunke-server/tests/test_smoke_script.py @@ -0,0 +1,164 @@ +"""Light unit tests for the end-to-end smoke script (Task 1.9). + +These DO NOT hit the network — they verify the testable surface of +:mod:`scripts.smoke_wechat`: + +* the module exposes an async ``main`` entrypoint, +* :func:`build_client_from_env` picks the right auth mode from env vars, +* :func:`_parse_cookie_env` accepts both raw-header and JSON cookie formats. + +The full :func:`main` flow is an integration check against the live CRM and is +exercised manually (see the module docstring of ``scripts.smoke_wechat``). +""" + +from __future__ import annotations + +import asyncio + +from app.session.yunke_client import YunkeClient +from scripts.smoke_wechat import _parse_cookie_env, build_client_from_env, main + +USERID = "a" * 32 +COMPANY = "8m0atd" + + +# --------------------------------------------------------------------------- +# Entrypoint shape +# --------------------------------------------------------------------------- + + +def test_main_is_async_entrypoint(): + """``main`` must be a coroutine function so ``asyncio.run(main())`` works.""" + assert asyncio.iscoroutinefunction(main) + + +# --------------------------------------------------------------------------- +# build_client_from_env — auth-mode selection +# --------------------------------------------------------------------------- + + +async def test_build_client_prefers_existing_session(monkeypatch): + """YK_USERID/YK_COMPANY/YK_COOKIE set => existing-session mode.""" + # Clear any conflicting env so the test is hermetic. + for k in ("YUNKE_USERNAME", "YUNKE_PASSWORD"): + monkeypatch.delenv(k, raising=False) + monkeypatch.setenv("YK_USERID", USERID) + monkeypatch.setenv("YK_COMPANY", COMPANY) + monkeypatch.setenv("YK_COOKIE", "JSESSIONID=deadbeef; foo=bar") + + client = build_client_from_env() + try: + assert isinstance(client, YunkeClient) + # Session mode => already logged in (login() must NOT be required). + assert client.logged_in + assert client.login_result is not None + assert client.login_result.userid == USERID + assert client.login_result.company == COMPANY + # Cookies from YK_COOKIE landed in the underlying httpx cookie jar. + assert client.client.cookies.get("JSESSIONID") == "deadbeef" + assert client.client.cookies.get("foo") == "bar" + # And login mode was NOT chosen (no creds attached). + assert client.username is None + assert client.password is None + finally: + # build_client_from_env constructs an owned httpx.AsyncClient; close + # it so the test does not leak the connection pool. + await client.aclose() + + +async def test_build_client_login_mode_when_only_creds(monkeypatch): + """YUNKE_USERNAME + YUNKE_PASSWORD (no session env) => credentialed client.""" + for k in ("YK_USERID", "YK_COMPANY", "YK_COOKIE"): + monkeypatch.delenv(k, raising=False) + monkeypatch.setenv("YUNKE_USERNAME", "alice") + monkeypatch.setenv("YUNKE_PASSWORD", "secret") + + client = build_client_from_env() + try: + assert isinstance(client, YunkeClient) + assert client.username == "alice" + assert client.password == "secret" + # Login mode => NOT logged in yet; the caller must call login(). + assert not client.logged_in + assert client.login_result is None + finally: + await client.aclose() + + +async def test_build_client_session_takes_priority_over_creds(monkeypatch): + """If both session and credential env are present, session mode wins.""" + monkeypatch.setenv("YK_USERID", USERID) + monkeypatch.setenv("YK_COMPANY", COMPANY) + monkeypatch.setenv("YUNKE_USERNAME", "alice") + monkeypatch.setenv("YUNKE_PASSWORD", "secret") + + client = build_client_from_env() + try: + assert isinstance(client, YunkeClient) + assert client.logged_in + # Session mode chosen; cred fields left unset. + assert client.username is None + finally: + await client.aclose() + + +def test_build_client_returns_none_when_no_env(monkeypatch): + """No auth env set => None (the caller prints usage and exits 2).""" + for k in ( + "YK_USERID", + "YK_COMPANY", + "YK_COOKIE", + "YUNKE_USERNAME", + "YUNKE_PASSWORD", + ): + monkeypatch.delenv(k, raising=False) + + assert build_client_from_env() is None + + +def test_build_client_raises_on_partial_session_env(monkeypatch): + """Only one of YK_USERID/YK_COMPANY set => ValueError (clear caller bug).""" + for k in ( + "YK_COMPANY", + "YK_COOKIE", + "YUNKE_USERNAME", + "YUNKE_PASSWORD", + ): + monkeypatch.delenv(k, raising=False) + monkeypatch.setenv("YK_USERID", USERID) + + import pytest + + with pytest.raises(ValueError, match="YK_USERID and YK_COMPANY"): + build_client_from_env() + + +# --------------------------------------------------------------------------- +# _parse_cookie_env — cookie string parsing +# --------------------------------------------------------------------------- + + +def test_parse_cookie_env_raw_header(): + assert _parse_cookie_env("JSESSIONID=abc; foo=bar") == { + "JSESSIONID": "abc", + "foo": "bar", + } + + +def test_parse_cookie_env_json(): + assert _parse_cookie_env('{"JSESSIONID": "abc", "x": "y"}') == { + "JSESSIONID": "abc", + "x": "y", + } + + +def test_parse_cookie_env_empty_returns_empty_dict(): + assert _parse_cookie_env("") == {} + assert _parse_cookie_env(" ") == {} + + +def test_parse_cookie_env_malformed_returns_empty_dict(): + # Non-JSON, no equals signs -> nothing parseable -> empty. + assert _parse_cookie_env("not a cookie") == {} + # Malformed JSON -> empty (no raise). + assert _parse_cookie_env('{bad json') == {}