Files
yunke/yunke-server/tests/test_scrape_whatsapp_messages.py
T
浅念 8a22fd9ff8 feat: whatsapp account/friend/message scrapers
Mirrors the WeChat scraper patterns (M3) using the shared paginator +
parser + repo building blocks. Three new async generators yield
snake_case dicts + verbatim raw_json, ready for the channel-generic
upsert_customer / upsert_messages repos.

- app/scraper/whatsapp.py (new): scrape_whatsapp_accounts /
  scrape_whatsapp_friends (page-number pagination over the
  data.overviews / data.totalCount envelope) and scrape_whatsapp_messages
  (backward-cursor pagination, JSON body {uid,friendUid,start,end} —
  NOT form-encoded like WeChat msg/chats). Drops type==-2 dividers;
  maps mediaName -> file_url as a deferred-download key.
- app/scraper/parsers.py: add normalize_whatsapp_type (0->text,
  1->image, -2->divider, else unknown; never raises) and a shared
  coerce_int_or_none scraper helper (mirrors repos._coerce_int so the
  scraper emits type_raw ints without a DB dependency).
- app/storage/repos.py: add upsert_whatsapp_account (PK uid) which
  also upserts the embedded Staff(user_id, name) via the existing
  upsert_staff_from_account helper; WhatsappAccount imported.

Tests (respx mocks, no network): 3 cassette fixtures + 3 scraper test
modules (accounts/friends/messages), parsers + repos extended with WA
cases. Divider drop, backward-cursor advance (end == E1-1), JSON body
asserts, embedded-staff upsert, and channel-generic Customer/Message
upserts all covered.

Suite: 182 -> 211 passing. ruff clean.

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-01 15:48:03 +08:00

209 lines
6.8 KiB
Python

"""Tests for the WhatsApp message scraper (M3).
Drives the real backward-cursor pagination + normalization logic through
``respx`` by wiring a real ``YunkeClient`` (primed with a fake session so the
login flow is skipped) against a mocked transport. Two pages are served:
* page 1 -> ``messages=[m1, m_datebar(-2), m2]``, ``end=E1``, ``hasLast=true``
* page 2 -> ``messages=[m3]``, ``end=E2``, ``hasLast=false``
The scraper must yield exactly ``[m1, m2, m3]`` (the ``-2`` divider dropped),
map fields per the live-capture names, set ``channel="whatsapp"``, send a
**JSON** body (``uid``/``friendUid``/``start``/``end``), AND the 2nd request
body's ``end`` must equal ``E1 - 1`` (proving the backward cursor advances by
one ms).
"""
from __future__ import annotations
import json as _json
from typing import Any
import httpx
import respx
from app.scraper.whatsapp import scrape_whatsapp_messages
from app.session.yunke_client import YunkeClient
from tests.cassettes.wa_messages_fixture import (
E1,
E2,
WA_MESSAGE_RECORDS_PAGE1,
WA_MESSAGE_RECORDS_PAGE2,
make_wa_messages_response,
)
BASE = "https://crm2.yunkecn.com"
WA_MESSAGES_URL = f"{BASE}/qqweb/WhatsAppUser/chatRecords"
USERID = "a" * 32
COMPANY = "8m0atd"
UID = "8613800000001@s.whatsapp.net"
FRIEND_UID = "8613800000111@s.whatsapp.net"
START = 1
END = 1784518000000 # the initial upper-bound cursor
def _primed_client() -> YunkeClient:
"""A YunkeClient already in a logged-in state (login flow skipped)."""
return YunkeClient(base_url=BASE, userid=USERID, company=COMPANY)
def _json_body(sent: httpx.Request) -> dict[str, Any]:
"""Parse the JSON body of an outgoing request into a dict."""
return _json.loads(sent.content.decode())
# ---------------------------------------------------------------------------
# Test A: end-to-end backward-cursor pagination + divider drop + field mapping
# ---------------------------------------------------------------------------
@respx.mock
async def test_scrape_whatsapp_messages_drops_divider_and_paginates_backward():
page1 = make_wa_messages_response(
WA_MESSAGE_RECORDS_PAGE1, end=E1, has_last=True, trace_id="trace-wa-1"
)
page2 = make_wa_messages_response(
WA_MESSAGE_RECORDS_PAGE2, end=E2, has_last=False, trace_id="trace-wa-2"
)
route = respx.post(WA_MESSAGES_URL).mock(
side_effect=[
httpx.Response(200, json=page1),
httpx.Response(200, json=page2),
]
)
client = _primed_client()
try:
results = [
r
async for r in scrape_whatsapp_messages(
client, UID, FRIEND_UID, START, END
)
]
finally:
await client.aclose()
# --- Exactly 2 calls: hasLast=false on page 2 means no 3rd fetch. ---
assert route.call_count == 2
# --- Exactly [m1, m2, m3] yielded (the -2 divider dropped). ---
assert [r["msg_id"] for r in results] == ["wa_1000001", "wa_1000002", "wa_999999"]
assert len(results) == 3
r1, r2, r3 = results
# m1 — text / incoming.
assert r1["msg_id"] == "wa_1000001"
assert r1["channel"] == "whatsapp"
assert r1["owner_account_id"] == UID
assert r1["talker_id"] == FRIEND_UID
assert r1["is_send"] == 0
assert r1["type"] == "text" # code 0 -> text (phase0 §3)
assert r1["type_raw"] == 0 # original raw code preserved
assert r1["text"] == "hello from whatsapp"
assert r1["timestamp"] == E1 + 600000
assert r1["sender_name"] == "Customer One"
assert r1["raw_json"] == WA_MESSAGE_RECORDS_PAGE1[0]
# m2 — image / outgoing / media filename (deferred download, not a URL).
assert r2["msg_id"] == "wa_1000002"
assert r2["is_send"] == 1
assert r2["type"] == "image" # code 1 -> image/media (phase0 §3)
assert r2["type_raw"] == 1
assert r2["text"] == "[image]"
assert r2["file_url"] == "8613800000001@s.whatsapp.net_wa_1000002.jpg"
assert r2["raw_json"] == WA_MESSAGE_RECORDS_PAGE1[2]
# m3 — text / incoming / quoted reply.
assert r3["msg_id"] == "wa_999999"
assert r3["type"] == "text"
assert r3["type_raw"] == 0
assert r3["text"] == "replying to an earlier message"
assert r3["refer_msg"] == {
"quotedId": "wa_888888",
"quotedMessage": "the earlier message text",
}
assert r3["raw_json"] == WA_MESSAGE_RECORDS_PAGE2[0]
# --- Backward cursor advance: 2nd request body's end == E1 - 1. ---
first_req = route.calls[0].request
second_req = route.calls[1].request
body1 = _json_body(first_req)
body2 = _json_body(second_req)
assert body1["end"] == END
assert body2["end"] == E1 - 1
# --- JSON body is actually sent (NOT form-encoded). ---
assert first_req.headers["content-type"] == "application/json"
assert body1["uid"] == UID
assert body1["friendUid"] == FRIEND_UID
assert body1["start"] == START
assert body1["end"] == END
# ``start`` is constant across pages (only ``end`` walks backward).
assert body2["start"] == START
assert body2["uid"] == UID
assert body2["friendUid"] == FRIEND_UID
# ---------------------------------------------------------------------------
# Test B: hasLast=false on first page -> stop after one call
# ---------------------------------------------------------------------------
@respx.mock
async def test_scrape_whatsapp_messages_stops_when_has_last_false_on_first_page():
"""A single page with hasLast=false: scraper yields the records and stops."""
page = make_wa_messages_response(
[WA_MESSAGE_RECORDS_PAGE2[0]], end=E2, has_last=False
)
route = respx.post(WA_MESSAGES_URL).mock(
return_value=httpx.Response(200, json=page)
)
client = _primed_client()
try:
results = [
r
async for r in scrape_whatsapp_messages(
client, UID, FRIEND_UID, START, END
)
]
finally:
await client.aclose()
assert [r["msg_id"] for r in results] == ["wa_999999"]
assert route.call_count == 1
# ---------------------------------------------------------------------------
# Test C: empty messages array on the first page -> stop without a 3rd fetch
# ---------------------------------------------------------------------------
@respx.mock
async def test_scrape_whatsapp_messages_empty_first_page_yields_nothing():
"""An empty messages array is a defensive stop signal — no rows, one call."""
page = make_wa_messages_response([], end=E1, has_last=True)
route = respx.post(WA_MESSAGES_URL).mock(
return_value=httpx.Response(200, json=page)
)
client = _primed_client()
try:
results = [
r
async for r in scrape_whatsapp_messages(
client, UID, FRIEND_UID, START, END
)
]
finally:
await client.aclose()
assert results == []
assert route.call_count == 1