Files
yunke/yunke-server/tests/test_repos.py
T
浅念 2b79e7a232 feat: department + wechat account/staff scrapers
M2 organization scrapers for yunke-server:

- scrape_departments: POST /usercenter/dept/tree. Single request (no
  pagination). Defensively unwraps the envelope (data is node list OR an
  object wrapping nodes under children), flattens BOTH nested and flat
  shapes, and reconstructs the materialized path by walking the pId chain
  (ancestor names joined with /). Maps id/pId/name -> dept_id/parent_id/name.

- scrape_wechat_accounts: POST /wechatweb/wxUser/list. Page-number
  pagination via the existing paginate_by_page helper. Envelope differs
  from friend_list (no basicResult wrapper): data.data/data.total. Maps
  wechatId/nickName/alias/phone/userId/userName/contactsNum/groupNum/
  sendNum/receiveNum/isOnJob/dbSize/onlineStatus/lastUpdate/version to
  WechatAccount columns; carries embedded staff pair for the repo layer.

Repos:
- upsert_department: idempotent upsert into Department (PK dept_id).
- upsert_staff: primitive idempotent staff upsert (PK user_id, is_on_job
  coerced to bool).
- upsert_wechat_account: now ALSO upserts the embedded Staff row via
  upsert_staff_from_account (best-effort; skips when account has no user_id).
- upsert_staff_from_account: refactored to delegate to upsert_staff and
  forward phone/sale_phone/is_on_job when present.

Tests (TDD red->green): test_scrape_departments (nested + flat + wrapped-data
+ empty), test_scrape_wechat_accounts (2-page pagination + body advance +
empty first page), test_repos extended (dept idempotency/root parent, direct
staff upsert + bool coercion, account-also-upserts-staff, re-upsert updates
both, no-user_id skips staff). Cassettes: dept_tree_fixture (same 4 nodes in
both nested + flat shapes), wxUser_fixture (4 capture-shaped accounts incl.
departed + orphan).

Suite: 97 -> 111 passing. ruff check clean.

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-01 13:46:23 +08:00

704 lines
25 KiB
Python

"""Tests for app.storage.repos — idempotent upserts + type coercion.
These tests build an isolated in-memory SQLite async engine (same pattern as
``tests/test_models.py``) and exercise the public upsert helpers in
``app.storage.repos``. They never touch the production MySQL engine.
Coverage:
* Idempotency: same PK twice -> one row, with non-PK fields updated.
* Type coercion: bool fields (``is_send``/``is_del``/``is_block``/``is_on_job``)
accept int / "1"-style str / bool; ``gender`` str -> int; ``None`` preserved.
* Batch dedup: 3 messages with one duplicate ``msg_id`` -> 2 rows in the table.
* Return values: ``upsert_customers`` / ``upsert_messages`` report counts.
* Embedded-staff best-effort upsert from a wechat-account record.
* Direct unit tests on the module-level ``_coerce_bool`` / ``_coerce_int``
helpers (kept module-level precisely so they can be exercised here).
"""
from __future__ import annotations
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from app.db.base import Base
from app.db.models import Customer, Department, Message, Staff, WechatAccount
from app.storage.repos import (
_coerce_bool,
_coerce_int,
upsert_customer,
upsert_customers,
upsert_department,
upsert_messages,
upsert_staff,
upsert_staff_from_account,
upsert_wechat_account,
)
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
async def _make_engine():
"""Create an in-memory async SQLite engine with all tables created."""
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
return engine
# ---------------------------------------------------------------------------
# Coercion helpers (module-level, exercised directly for unit coverage)
# ---------------------------------------------------------------------------
def test_coerce_bool_handles_documented_cases():
"""bool fields accept None / bool / int / '1'-style str."""
assert _coerce_bool(None) is None
assert _coerce_bool(True) is True
assert _coerce_bool(False) is False
assert _coerce_bool(1) is True
assert _coerce_bool(0) is False
assert _coerce_bool("1") is True
assert _coerce_bool("0") is False
def test_coerce_int_handles_documented_cases():
"""gender / type-style fields: str -> int; None preserved; 0 preserved."""
assert _coerce_int(None) is None
assert _coerce_int(2) == 2
assert _coerce_int("2") == 2
assert _coerce_int(0) == 0
assert _coerce_int("0") == 0
def test_coerce_int_returns_none_for_non_numeric_string():
"""A non-numeric string (e.g. a label) -> None.
Integer columns (``type_raw``, ``gender``, ...) cannot store a
non-numeric value; the coercion collapses it to ``None``.
"""
assert _coerce_int("text") is None
assert _coerce_int("unknown") is None
# ---------------------------------------------------------------------------
# Customer upsert
# ---------------------------------------------------------------------------
async def test_upsert_customer_is_idempotent_and_updates_version():
"""Upsert same friend_id twice with different version -> 1 row, version updated."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
rec1 = {
"friend_id": "wxid_abc",
"channel": "wechat",
"version": "v1",
"friend_nickname": "Alice",
}
rec2 = {
"friend_id": "wxid_abc",
"channel": "wechat",
"version": "v2",
"friend_nickname": "Alice",
}
await upsert_customer(session, rec1)
await session.commit()
await upsert_customer(session, rec2)
await session.commit()
# ``populate_existing=True`` bypasses the identity-map cache so the
# verify query reflects the second upsert's writes (with
# ``expire_on_commit=False`` the loaded instance is otherwise stale).
rows = (
await session.execute(
sa.select(Customer).execution_options(populate_existing=True)
)
).scalars().all()
assert len(rows) == 1
assert rows[0].friend_id == "wxid_abc"
assert rows[0].version == "v2"
assert rows[0].friend_nickname == "Alice"
finally:
await engine.dispose()
async def test_upsert_customer_coerces_gender_and_bool():
"""gender '2' -> int 2; is_block '1' -> True."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
rec = {
"friend_id": "wxid_g1",
"channel": "wechat",
"gender": "2",
"is_block": "1",
}
await upsert_customer(session, rec)
await session.commit()
fetched = (
await session.execute(
sa.select(Customer).where(Customer.friend_id == "wxid_g1")
)
).scalar_one()
assert fetched.gender == 2
assert fetched.is_block is True
finally:
await engine.dispose()
async def test_upsert_customer_preserves_none_for_bool_and_gender():
"""Missing / None bool & gender stay None (no spurious coercion)."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
rec = {
"friend_id": "wxid_n",
"channel": "wechat",
"gender": None,
"is_block": None,
}
await upsert_customer(session, rec)
await session.commit()
fetched = (
await session.execute(
sa.select(Customer).where(Customer.friend_id == "wxid_n")
)
).scalar_one()
assert fetched.gender is None
assert fetched.is_block is None
finally:
await engine.dispose()
async def test_upsert_customer_keeps_raw_json():
"""raw_json is preserved verbatim (the whole input record)."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
rec = {
"friend_id": "wxid_r",
"channel": "wechat",
"friendWxNickname": "upstream-name", # a non-column key
"nested": {"k": "v"},
}
await upsert_customer(session, rec)
await session.commit()
fetched = (
await session.execute(
sa.select(Customer).where(Customer.friend_id == "wxid_r")
)
).scalar_one()
assert fetched.raw_json == rec
finally:
await engine.dispose()
async def test_upsert_customers_batch_returns_count():
"""Batch upsert returns the number of records processed."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
recs = [
{"friend_id": "wxid_b1", "channel": "wechat", "version": "v1"},
{"friend_id": "wxid_b2", "channel": "wechat", "version": "v1"},
]
count = await upsert_customers(session, recs)
await session.commit()
assert count == 2
rows = (await session.execute(sa.select(Customer))).scalars().all()
assert {r.friend_id for r in rows} == {"wxid_b1", "wxid_b2"}
finally:
await engine.dispose()
# ---------------------------------------------------------------------------
# Message upsert
# ---------------------------------------------------------------------------
async def test_upsert_messages_coerces_types():
"""is_send '1' -> True; is_del 0 -> False; type label stored as string;
type_raw coerced to int; same msg_id twice -> 1 row."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
rec = {
"msg_id": "m_coerce",
"channel": "wechat",
"is_send": "1",
"is_del": 0,
"type": "text", # normalized label — stored as-is (no coercion)
"type_raw": "1", # raw code arrives as str -> coerced to int
"timestamp": 1_700_000_000_000,
}
await upsert_messages(session, [rec])
await session.commit()
fetched = (
await session.execute(
sa.select(Message).where(Message.msg_id == "m_coerce")
)
).scalar_one()
assert fetched.is_send is True
assert fetched.is_del is False
assert fetched.type == "text" # string label preserved
assert fetched.type_raw == 1 # raw code coerced to int
assert fetched.timestamp == 1_700_000_000_000
# Upsert same msg_id again with different text + new raw code ->
# row updated, not duplicated.
rec2 = dict(rec, text="updated", type="unknown", type_raw=3)
await upsert_messages(session, [rec2])
await session.commit()
# ``populate_existing=True`` bypasses the identity-map cache so we
# observe the second upsert's values.
rows = (
await session.execute(
sa.select(Message)
.where(Message.msg_id == "m_coerce")
.execution_options(populate_existing=True)
)
).scalars().all()
assert len(rows) == 1
assert rows[0].text == "updated"
assert rows[0].type == "unknown"
assert rows[0].type_raw == 3
finally:
await engine.dispose()
async def test_upsert_messages_preserves_none_type_and_type_raw():
"""Missing / None type & type_raw stay None (no spurious coercion)."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
rec = {
"msg_id": "m_none",
"channel": "wechat",
"type": None,
"type_raw": None,
}
await upsert_messages(session, [rec])
await session.commit()
fetched = (
await session.execute(
sa.select(Message).where(Message.msg_id == "m_none")
)
).scalar_one()
assert fetched.type is None
assert fetched.type_raw is None
finally:
await engine.dispose()
async def test_upsert_messages_dedupes_duplicate_msg_id():
"""3 messages with 1 duplicate msg_id -> 2 rows; dup row updated."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
recs = [
{"msg_id": "d1", "channel": "wechat", "text": "a"},
{"msg_id": "d2", "channel": "wechat", "text": "b"},
{"msg_id": "d1", "channel": "wechat", "text": "a-updated"},
]
count = await upsert_messages(session, recs)
await session.commit()
rows = (
await session.execute(
sa.select(Message).execution_options(populate_existing=True)
)
).scalars().all()
assert len(rows) == 2
assert {r.msg_id for r in rows} == {"d1", "d2"}
# The duplicate's later text wins.
d1 = next(r for r in rows if r.msg_id == "d1")
assert d1.text == "a-updated"
# Return value reflects records processed (input count).
assert count == 3
finally:
await engine.dispose()
async def test_upsert_messages_empty_list_returns_zero():
"""An empty batch is a no-op and reports 0."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
count = await upsert_messages(session, [])
assert count == 0
finally:
await engine.dispose()
# ---------------------------------------------------------------------------
# WechatAccount + embedded Staff upsert (best-effort)
# ---------------------------------------------------------------------------
async def test_upsert_wechat_account_coerces_is_on_job():
"""WechatAccount upsert stores wx_id PK and coerces is_on_job."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
rec = {
"wx_id": "wxid_staff",
"nick_name": "Alice",
"user_id": "u1",
"user_name": "Alice Wang",
"is_on_job": 1,
"contacts_num": 100,
}
await upsert_wechat_account(session, rec)
await session.commit()
fetched = (
await session.execute(
sa.select(WechatAccount).where(
WechatAccount.wx_id == "wxid_staff"
)
)
).scalar_one()
assert fetched.nick_name == "Alice"
assert fetched.is_on_job is True
assert fetched.contacts_num == 100
finally:
await engine.dispose()
async def test_upsert_staff_from_account_extracts_user_id_and_name():
"""Embedded staff (user_id + user_name) upsert from an account record."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
rec = {
"wx_id": "wxid_staff",
"user_id": "u1",
"user_name": "Alice Wang",
"is_on_job": 1,
}
await upsert_staff_from_account(session, rec)
await session.commit()
fetched = (
await session.execute(
sa.select(Staff).where(Staff.user_id == "u1")
)
).scalar_one()
assert fetched.name == "Alice Wang"
assert fetched.is_on_job is True
finally:
await engine.dispose()
async def test_upsert_staff_from_account_skips_when_no_user_id():
"""No user_id -> nothing inserted (best-effort, no raise)."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
rec = {"wx_id": "wxid_no_user", "nick_name": "anon"}
# Must not raise.
await upsert_staff_from_account(session, rec)
await session.commit()
rows = (await session.execute(sa.select(Staff))).scalars().all()
assert rows == []
finally:
await engine.dispose()
# ---------------------------------------------------------------------------
# Department upsert (Task 2.0)
# ---------------------------------------------------------------------------
async def test_upsert_department_is_idempotent_and_updates_path():
"""Same dept_id twice with different path -> 1 row, path refreshed."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
rec1 = {
"dept_id": "1.1",
"parent_id": "1",
"name": "East",
"path": "Sales/East",
}
rec2 = {
"dept_id": "1.1",
"parent_id": "1",
"name": "East",
# Re-scrape reconstructs the same path; verify it is overwritten
# rather than producing a duplicate row.
"path": "Sales/East",
"raw_json": {"id": "1.1", "pId": "1", "name": "East"},
}
await upsert_department(session, rec1)
await session.commit()
await upsert_department(session, rec2)
await session.commit()
rows = (
await session.execute(
sa.select(Department)
.where(Department.dept_id == "1.1")
.execution_options(populate_existing=True)
)
).scalars().all()
assert len(rows) == 1
assert rows[0].dept_id == "1.1"
assert rows[0].parent_id == "1"
assert rows[0].name == "East"
assert rows[0].path == "Sales/East"
# raw_json from the second upsert wins (last write).
assert rows[0].raw_json == rec2["raw_json"]
finally:
await engine.dispose()
async def test_upsert_department_preserves_none_parent_for_root():
"""A root dept (parent_id None) is stored as-is; no spurious coercion."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
rec = {
"dept_id": "1",
"parent_id": None,
"name": "Sales",
"path": "Sales",
}
await upsert_department(session, rec)
await session.commit()
fetched = (
await session.execute(
sa.select(Department).where(Department.dept_id == "1")
)
).scalar_one()
assert fetched.parent_id is None
assert fetched.name == "Sales"
assert fetched.path == "Sales"
finally:
await engine.dispose()
# ---------------------------------------------------------------------------
# upsert_staff — direct staff-shaped dict upsert (Task 2.1)
# ---------------------------------------------------------------------------
async def test_upsert_staff_is_idempotent_and_updates_name():
"""Same user_id twice with different name -> 1 row, name refreshed."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
rec1 = {"user_id": "u_new", "name": "Old Name"}
rec2 = {"user_id": "u_new", "name": "New Name"}
await upsert_staff(session, rec1)
await session.commit()
await upsert_staff(session, rec2)
await session.commit()
rows = (
await session.execute(
sa.select(Staff)
.where(Staff.user_id == "u_new")
.execution_options(populate_existing=True)
)
).scalars().all()
assert len(rows) == 1
assert rows[0].name == "New Name"
finally:
await engine.dispose()
async def test_upsert_staff_coerces_is_on_job():
"""is_on_job '1' -> True; '0' -> False; None preserved."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
await upsert_staff(
session, {"user_id": "u_b1", "is_on_job": "1"}
)
await upsert_staff(
session, {"user_id": "u_b0", "is_on_job": "0"}
)
await upsert_staff(
session, {"user_id": "u_bn", "is_on_job": None}
)
await session.commit()
on = (
await session.execute(
sa.select(Staff).where(Staff.user_id == "u_b1")
)
).scalar_one()
off = (
await session.execute(
sa.select(Staff).where(Staff.user_id == "u_b0")
)
).scalar_one()
none_ = (
await session.execute(
sa.select(Staff).where(Staff.user_id == "u_bn")
)
).scalar_one()
assert on.is_on_job is True
assert off.is_on_job is False
assert none_.is_on_job is None
finally:
await engine.dispose()
# ---------------------------------------------------------------------------
# upsert_wechat_account now ALSO upserts the embedded Staff (Task 2.1)
# ---------------------------------------------------------------------------
async def test_upsert_wechat_account_also_upserts_embedded_staff():
"""Account upsert creates BOTH the WechatAccount row and the Staff row."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
rec = {
"wx_id": "wxid_dup",
"nick_name": "Dup WeChat",
"user_id": "u_dup",
"user_name": "Dup Real Name",
"is_on_job": 1,
"phone": "13800000099",
}
await upsert_wechat_account(session, rec)
await session.commit()
acct = (
await session.execute(
sa.select(WechatAccount).where(WechatAccount.wx_id == "wxid_dup")
)
).scalar_one()
assert acct.user_id == "u_dup"
assert acct.user_name == "Dup Real Name"
staff = (
await session.execute(
sa.select(Staff).where(Staff.user_id == "u_dup")
)
).scalar_one()
assert staff.name == "Dup Real Name"
assert staff.is_on_job is True
finally:
await engine.dispose()
async def test_upsert_wechat_account_reupsert_updates_embedded_staff():
"""Re-upserting the account refreshes BOTH account + embedded staff rows."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
rec1 = {
"wx_id": "wxid_re",
"nick_name": "Old Nick",
"user_id": "u_re",
"user_name": "Old Name",
"is_on_job": 1,
}
rec2 = {
"wx_id": "wxid_re",
"nick_name": "New Nick",
"user_id": "u_re",
"user_name": "New Name",
"is_on_job": 0, # staff departed since last scrape
}
await upsert_wechat_account(session, rec1)
await session.commit()
await upsert_wechat_account(session, rec2)
await session.commit()
# One account row, refreshed.
acct_rows = (
await session.execute(
sa.select(WechatAccount)
.where(WechatAccount.wx_id == "wxid_re")
.execution_options(populate_existing=True)
)
).scalars().all()
assert len(acct_rows) == 1
assert acct_rows[0].nick_name == "New Nick"
# One staff row, refreshed (not duplicated).
staff_rows = (
await session.execute(
sa.select(Staff)
.where(Staff.user_id == "u_re")
.execution_options(populate_existing=True)
)
).scalars().all()
assert len(staff_rows) == 1
assert staff_rows[0].name == "New Name"
assert staff_rows[0].is_on_job is False
finally:
await engine.dispose()
async def test_upsert_wechat_account_without_user_id_skips_staff():
"""Account with no user_id is written; embedded-staff upsert is skipped."""
engine = await _make_engine()
try:
SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
async with SessionLocal() as session:
rec = {
"wx_id": "wxid_no_staff",
"nick_name": "Orphan WeChat",
# no user_id / user_name
}
await upsert_wechat_account(session, rec)
await session.commit()
# Account row written.
acct = (
await session.execute(
sa.select(WechatAccount).where(
WechatAccount.wx_id == "wxid_no_staff"
)
)
).scalar_one()
assert acct.nick_name == "Orphan WeChat"
# No staff rows at all.
staff_rows = (
await session.execute(sa.select(Staff))
).scalars().all()
assert staff_rows == []
finally:
await engine.dispose()