浅念
854efd3119
Adds the storage layer that bridges scrapers (which yield normalized dicts)
to the DB and the local media store.
app/storage/repos.py
* upsert_customer / upsert_customers / upsert_messages / upsert_wechat_account
/ upsert_staff_from_account, all taking an AsyncSession + scraper-shaped dicts.
* Dialect-aware Core upsert: SQLite on_conflict_do_update for tests,
MySQL on_duplicate_key_update for prod; dialect detected via session.bind.name.
_build_upsert_stmt(model, values, dialect) is the reusable helper.
* Type coercion: _coerce_bool / _coerce_int (module-level, unit-tested) map
is_send/is_del/is_block/is_on_job int-as-bool to bool and gender/type
numeric strings to int; None preserved; raw_json kept verbatim.
* For M1 the Message model maps to a single physical `message` table;
per-month sharding (message_table_name helper) is a LATER task (YAGNI).
app/storage/object_store.py
* Local-filesystem ObjectStore with put_bytes / put_stream / put_by_hash /
exists / get_path, all dispatched via asyncio.to_thread so disk I/O is
genuinely async.
* put_by_hash computes sha256 and writes only if absent (content-addressed
dedup); layout is base_dir/<first 2 hash chars>/<hash><ext>.
* Default base dir is Settings.media_dir; the interface is the seam a
future MinIO/S3 backend will swap in behind.
tests: 24 new tests (in-memory SQLite for repos; tmp_path for object_store)
covering idempotency (same PK twice -> 1 row, updated), type coercion, batch
dedup (3 messages / 1 dup msg_id -> 2 rows), and hash dedup (corrupt-file
probe proves second put_by_hash does not rewrite). ruff clean; full suite
67/67 green.
Co-Authored-By: Claude <noreply@anthropic.com>
398 lines
14 KiB
Python
398 lines
14 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, Message, Staff, WechatAccount
|
|
from app.storage.repos import (
|
|
_coerce_bool,
|
|
_coerce_int,
|
|
upsert_customer,
|
|
upsert_customers,
|
|
upsert_messages,
|
|
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. normalized 'type' label) -> None.
|
|
|
|
The column is INTEGER; a non-numeric value cannot be stored. The raw value
|
|
remains recoverable from ``raw_json``.
|
|
"""
|
|
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; 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": "1",
|
|
"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 == 1
|
|
assert fetched.timestamp == 1_700_000_000_000
|
|
|
|
# Upsert same msg_id again with different text -> row updated, not
|
|
# duplicated.
|
|
rec2 = dict(rec, text="updated")
|
|
await upsert_messages(session, [rec2])
|
|
await session.commit()
|
|
|
|
# ``populate_existing=True`` bypasses the identity-map cache so we
|
|
# observe the second upsert's text value.
|
|
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"
|
|
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()
|