xinyouhe/yunke
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: idempotent upsert repos + object store

Browse Source 
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>
This commit is contained in:
浅念
2026-08-01 11:55:50 +08:00
parent 903286f379
commit 854efd3119
5 changed files with 1082 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
yunke-server/app/storage/__init__.py
+34
View File
@@ -0,0 +1,34 @@
"""Storage layer for the yunke data-mirror system.
Bridges the scrapers (``app.scraper.*``, which yield normalized dicts) to the
database (``app.db.models.*``) and the local media store.
Public surface (M1 / Task 1.7):
* :mod:`app.storage.repos` — idempotent Core upsert helpers that take an
``AsyncSession`` plus scraper-shaped dicts and write to the target table.
The upsert is dialect-aware (SQLite for tests, MySQL for prod) so the same
call path works in both environments.
* :mod:`app.storage.object_store` — a local-filesystem object store with
sha256-content-addressed dedup. The interface (``put_bytes`` / ``put_stream``
/ ``put_by_hash`` / ``exists`` / ``get_path``) is the seam a future MinIO/S3
backend will swap in behind.
"""
from app.storage.object_store import ObjectStore
from app.storage.repos import (
upsert_customer,
upsert_customers,
upsert_messages,
upsert_staff_from_account,
upsert_wechat_account,
)
__all__ = [
"ObjectStore",
"upsert_customer",
"upsert_customers",
"upsert_messages",
"upsert_staff_from_account",
"upsert_wechat_account",
]
yunke-server/app/storage/object_store.py
+155
View File
@@ -0,0 +1,155 @@
"""Local-filesystem object store with sha256 content-addressed dedup.
This is the M1 storage backend for media blobs (image / voice / video / file
downloads referenced by :class:`app.db.models.Media`). The interface is the
seam a future MinIO/S3 backend will swap in behind:
* :meth:`ObjectStore.put_bytes` — write bytes under an opaque key.
* :meth:`ObjectStore.put_stream` — write an async-iterable stream under a key.
* :meth:`ObjectStore.put_by_hash` — content-addressed put with sha256 dedup:
computes the hash, and only writes when no file with that hash exists. The
path layout is ``base_dir/<first 2 hash chars>/<hash><ext>`` (a simple
fan-out directory so no single directory holds millions of entries).
* :meth:`ObjectStore.exists` / :meth:`get_path` — predicate / pure-path helpers.
All disk I/O is dispatched via :func:`asyncio.to_thread`, so the methods are
genuinely async and safe to call from a FastAPI handler / scraper loop without
blocking the event loop on a slow disk.
The default base directory is :attr:`Settings.media_dir`.
"""
from __future__ import annotations
import asyncio
import hashlib
from collections.abc import AsyncIterator
from pathlib import Path
from app.config import get_settings
__all__ = ["ObjectStore"]
class ObjectStore:
"""Async local-filesystem object store with sha256-based dedup."""
def __init__(self, base_dir: Path | None = None) -> None:
"""Create a store rooted at ``base_dir``.
Defaults to :attr:`Settings.media_dir` when omitted so callers can
``ObjectStore()`` and land in the configured media location.
"""
self.base_dir: Path = Path(base_dir) if base_dir is not None else get_settings().media_dir
# ------------------------------------------------------------------
# Path helpers (pure, no I/O)
# ------------------------------------------------------------------
def get_path(self, key: str) -> Path:
"""Return ``base_dir / <key>`` (a pure-path computation; no I/O).
``key`` may contain ``/`` separators (e.g. ``"foo/bar.txt"``); they
become real subdirectories on disk. The key is treated as opaque — no
hash is computed.
"""
return self.base_dir / key
def _hash_path(self, sha256_hex: str, ext: str) -> Path:
"""Return ``base_dir / <first 2 hash chars> / <hash><ext>``."""
if ext and not ext.startswith("."):
ext = "." + ext
return self.base_dir / sha256_hex[:2] / f"{sha256_hex}{ext}"
# ------------------------------------------------------------------
# Writes
# ------------------------------------------------------------------
async def put_bytes(self, key: str, data: bytes) -> Path:
"""Write ``data`` under ``key`` (overwrites if present).
Returns the resolved path. Always writes; for content-addressed dedup
use :meth:`put_by_hash`.
"""
def _write() -> Path:
path = self.get_path(key)
path.parent.mkdir(parents=True, exist_ok=True)
path.write_bytes(data)
return path
return await asyncio.to_thread(_write)
async def put_stream(
self, key: str, stream: AsyncIterator[bytes]
) -> Path:
"""Concatenate an async-byte-stream under ``key`` (overwrites).
Chunks are consumed in the event-loop thread; the file open + each
write is dispatched to a worker thread to avoid blocking on disk.
"""
path = self.get_path(key)
def _mkdir() -> None:
path.parent.mkdir(parents=True, exist_ok=True)
await asyncio.to_thread(_mkdir)
def _write_chunk(buf: list[bytes], f) -> None:
for chunk in buf:
f.write(chunk)
buf.clear()
# Stream chunks in the event loop but push each disk write to a
# thread so we never block on slow storage mid-iteration.
def _open_and_prep():
return open(path, "wb")
f = await asyncio.to_thread(_open_and_prep)
try:
buffer: list[bytes] = []
async for chunk in stream:
buffer.append(chunk)
# Flush once we have a meaningful amount buffered.
if sum(len(c) for c in buffer) >= 65536:
await asyncio.to_thread(_write_chunk, buffer, f)
if buffer:
await asyncio.to_thread(_write_chunk, buffer, f)
finally:
await asyncio.to_thread(f.close)
return path
async def put_by_hash(
self, data: bytes, *, ext: str = ""
) -> tuple[str, Path]:
"""Content-addressed put with sha256 dedup.
Computes ``sha256(data)``, derives ``base_dir/<2 chars>/<hash><ext>``,
and writes the blob **only if no file with that path already exists**
(deduplication across messages that reference the same blob via md5).
Returns ``(sha256_hex, path)``. On a hit the existing file is left
untouched (no rewrite), so dedup is observable.
"""
sha256_hex = hashlib.sha256(data).hexdigest()
path = self._hash_path(sha256_hex, ext)
def _write_if_missing() -> Path:
if path.exists():
return path
path.parent.mkdir(parents=True, exist_ok=True)
path.write_bytes(data)
return path
await asyncio.to_thread(_write_if_missing)
return sha256_hex, path
# ------------------------------------------------------------------
# Reads / predicates
# ------------------------------------------------------------------
async def exists(self, key: str) -> bool:
"""True iff ``key`` exists on disk under ``base_dir``."""
return await asyncio.to_thread(self.get_path(key).exists)
yunke-server/app/storage/repos.py
+340
View File
@@ -0,0 +1,340 @@
"""Idempotent upsert repositories bridging scraper dicts to ORM tables.
Every public helper takes an ``AsyncSession`` plus one-or-more scraper-shaped
dicts (the snake_case records that ``app.scraper.*`` yields) and writes them
to the corresponding table via **dialect-aware Core upsert**, so the same call
path is correct on:
* **in-memory SQLite** — used by the test suite (``sqlite+aiosqlite``); uses
``sqlalchemy.dialects.sqlite.insert(...).on_conflict_do_update(...)``, and
* **MySQL** — production (``mysql+asyncmy``); uses
``sqlalchemy.dialects.mysql.insert(...).on_duplicate_key_update(...)``.
The dialect is detected at write time via ``session.bind.name``, so callers
do not need to know which backend they target.
Type coercion
-------------
Upstream CRM values are not always typed the way the columns want them. The
following coercions are applied when building column values from the input
dict (see :func:`_coerce_bool`, :func:`_coerce_int`):
* Booleans (``is_block`` / ``is_send`` / ``is_del`` / ``is_on_job``):
``1`` / ``"1"`` / ``True`` -> ``True``; ``0`` / ``"0"`` / ``False`` ->
``False``; ``None`` preserved.
* ``gender``: ``"0"`` / ``"1"`` / ``"2"`` (str) -> ``int`` when not ``None``.
* ``type`` (Message): coerced to ``int``; a non-numeric string (e.g. a
normalized type label like ``"text"``) collapses to ``None`` — the raw code
is still recoverable from ``raw_json``.
* Timestamps (``last_chat_time`` / ``timestamp`` / ``last_update`` / ...):
passed through as-is (ms-epoch ``int`` or ``None``).
* ``raw_json``: the original record dict, preserved verbatim. When the input
already carries a ``raw_json`` key (scraper output does), that value wins;
otherwise the input dict itself is stored.
Sharding note
-------------
``app.db.models.message_table_name`` exists for the future per-month sharded
``message_<channel>_<YYYYMM>`` layout. For M1, the :class:`Message` model maps
to a single physical ``message`` table and :func:`upsert_messages` writes to
that one table. Multi-table routing is a LATER task (YAGNI for M1).
"""
from __future__ import annotations
from typing import Any
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import DeclarativeBase
from app.db.models import Customer, Message, Staff, WechatAccount
__all__ = [
"upsert_customer",
"upsert_customers",
"upsert_messages",
"upsert_staff_from_account",
"upsert_wechat_account",
]
# ---------------------------------------------------------------------------
# Type coercion helpers (module-level so they are unit-testable)
# ---------------------------------------------------------------------------
# Columns whose source values are int-as-bool and must land as actual bool.
_BOOL_FIELDS: frozenset[str] = frozenset(
{"is_block", "is_send", "is_del", "is_on_job", "is_upload"}
)
# Columns whose source values may arrive as numeric strings.
_INT_FIELDS: frozenset[str] = frozenset(
{
"gender",
"type",
"pic_type",
"friend_state",
"delete_flag",
"direction",
"incoming_call",
"call_status",
"qc_status",
}
)
def _coerce_bool(value: Any) -> bool | None:
"""Coerce int/str/bool CRM bool fields to a real bool.
``None`` is preserved. ``1``/``"1"``/``True`` -> ``True``;
``0``/``"0"``/``False`` -> ``False``. Common textual forms
(``"true"``/``"false"``/``"yes"``/``"no"``) are also accepted; any other
non-empty string falls back to truthiness so a stray value never raises.
"""
if value is None:
return None
if isinstance(value, bool):
return value
if isinstance(value, int):
return value != 0
if isinstance(value, str):
s = value.strip().lower()
if s in ("1", "true", "yes", "y", "t"):
return True
if s in ("0", "false", "no", "n", "f", ""):
return False
# Last-resort: try numeric compare, else truthiness of the raw string.
try:
return int(s) != 0
except ValueError:
return bool(s)
# Floats / other: Python truthiness.
return bool(value)
def _coerce_int(value: Any) -> int | None:
"""Coerce a possibly-string numeric value to ``int``; ``None`` preserved.
``"2"`` -> ``2``; ``2`` -> ``2``; ``0``/``"0"`` -> ``0``; ``None`` ->
``None``. A non-numeric string (e.g. ``"text"``) returns ``None`` — the
caller (Message.type) treats that as "no value", and the raw code is still
recoverable from ``raw_json``.
"""
if value is None:
return None
if isinstance(value, bool):
return int(value)
if isinstance(value, int):
return value
if isinstance(value, float):
return int(value)
if isinstance(value, str):
s = value.strip()
if not s:
return None
try:
return int(s)
except ValueError:
return None
return None
# ---------------------------------------------------------------------------
# Dialect-aware Core upsert
# ---------------------------------------------------------------------------
def _pk_col_names(model: type[DeclarativeBase]) -> list[str]:
"""Return the names of the model's primary-key columns."""
return [c.name for c in model.__table__.primary_key.columns]
def _build_column_values(
model: type[DeclarativeBase], rec: dict[str, Any]
) -> dict[str, Any]:
"""Pick rec keys that match model columns, applying type coercion.
* Only keys that correspond to real columns on ``model`` are kept (extra
upstream keys are dropped, since they would cause Core insert() to raise
on unknown columns).
* Boolean / int-as-str fields are coerced via :func:`_coerce_bool` /
:func:`_coerce_int`.
* ``raw_json`` defaults to the entire input dict when not explicitly
provided, so the verbatim upstream record is always recoverable.
"""
col_names = {c.name for c in model.__table__.columns}
values: dict[str, Any] = {}
for col_name in col_names:
if col_name in rec:
v = rec[col_name]
if col_name in _BOOL_FIELDS:
v = _coerce_bool(v)
elif col_name in _INT_FIELDS:
v = _coerce_int(v)
values[col_name] = v
# Preserve the verbatim source record. If the scraper already set raw_json
# (it does), that wins; otherwise the input dict itself is the source of
# truth.
if "raw_json" in col_names and "raw_json" not in values:
values["raw_json"] = rec
return values
def _build_upsert_stmt(
model: type[DeclarativeBase], values: dict[str, Any], dialect_name: str
) -> sa.Insert:
"""Build a dialect-specific idempotent Core upsert for one row.
On conflict the statement updates **every non-PK column** from the
incoming row, so re-running a scrape yields one final row per PK with the
freshest values (counts, version, names, raw_json, updated_at, ...).
Parameters
----------
model:
The ORM model class to upsert into.
values:
Column-name -> value dict (output of :func:`_build_column_values`).
dialect_name:
One of ``"sqlite"`` / ``"mysql"`` (the two backends this project
targets). Any other dialect raises ``ValueError``.
"""
pk_cols = _pk_col_names(model)
# Non-PK columns to refresh on conflict (volatile fields).
update_col_names = [c.name for c in model.__table__.columns if c.name not in pk_cols]
# NB: access ``excluded`` / ``inserted`` on the BARE insert (before
# ``.values(...)``) and reuse that accessor when building the SET dict.
# If the accessor is taken on the post-``.values()`` statement, the
# ``excluded.<col>`` references do not bind to the supplied values at
# compile time and the ON CONFLICT DO UPDATE silently fails to refresh
# the row (a SQLAlchemy quirk; verified on 2.0.x).
if dialect_name == "sqlite":
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
bare = sqlite_insert(model)
excluded = bare.excluded
set_ = {name: excluded[name] for name in update_col_names}
return bare.values(**values).on_conflict_do_update(
index_elements=pk_cols, set_=set_
)
if dialect_name == "mysql":
from sqlalchemy.dialects.mysql import insert as mysql_insert
bare = mysql_insert(model)
inserted = bare.inserted
update_dict = {name: inserted[name] for name in update_col_names}
return bare.values(**values).on_duplicate_key_update(**update_dict)
raise ValueError(
f"Unsupported dialect for idempotent upsert: {dialect_name!r}. "
"Only 'sqlite' (tests) and 'mysql' (prod) are supported."
)
async def _dialect_name(session: AsyncSession) -> str:
"""Return the dialect name of the session's bind (e.g. 'sqlite'/'mysql').
Uses ``session.bind.name`` (the AsyncEngine); falls back to
``session.get_bind().dialect.name`` if ``bind`` is not set on the session.
"""
bind = session.bind
if bind is not None and hasattr(bind, "name"):
return bind.name # type: ignore[no-any-return]
return session.get_bind().dialect.name
async def _upsert_one(session: AsyncSession, model: type[DeclarativeBase], rec: dict[str, Any]) -> None:
"""Upsert a single record into ``model``."""
values = _build_column_values(model, rec)
# If the record carries no values for any column (e.g. empty dict), there
# is nothing to write — skip the round-trip. PK-only rows still go through
# because their PK is in `values`.
if not values:
return
dialect = await _dialect_name(session)
stmt = _build_upsert_stmt(model, values, dialect)
await session.execute(stmt)
# ---------------------------------------------------------------------------
# Public upsert helpers
# ---------------------------------------------------------------------------
async def upsert_customer(session: AsyncSession, rec: dict[str, Any]) -> None:
"""Idempotently upsert one :class:`Customer` (PK = ``friend_id``).
Same ``friend_id`` twice -> one row, with non-PK fields refreshed from the
later record.
"""
await _upsert_one(session, Customer, rec)
async def upsert_customers(session: AsyncSession, recs: list[dict[str, Any]]) -> int:
"""Batch upsert for :class:`Customer`. Returns the number of records processed.
Idempotent: duplicate ``friend_id`` values within the batch collapse to a
single row (the last write wins) at the DB layer.
"""
for rec in recs:
await _upsert_one(session, Customer, rec)
return len(recs)
async def upsert_messages(session: AsyncSession, recs: list[dict[str, Any]]) -> int:
"""Batch idempotent upsert for :class:`Message` (PK = ``msg_id``).
Applies the documented coercions (``is_send``/``is_del`` -> bool, ``type``
-> int). Returns the number of records processed (input length); the DB
collapses duplicate ``msg_id`` values to one row.
.. note::
For M1 the ``Message`` model maps to a single physical ``message``
table. Real per-month sharding across
``message_<channel>_<YYYYMM>`` tables (see
:func:`app.db.models.message_table_name`) is a LATER task; do not
implement multi-table routing here.
"""
for rec in recs:
await _upsert_one(session, Message, rec)
return len(recs)
# ---------------------------------------------------------------------------
# WechatAccount + embedded Staff (best-effort)
# ---------------------------------------------------------------------------
async def upsert_wechat_account(session: AsyncSession, rec: dict[str, Any]) -> None:
"""Best-effort idempotent upsert for :class:`WechatAccount` (PK = ``wx_id``).
The companion :func:`upsert_staff_from_account` covers the embedded staff
row (``user_id`` + ``user_name``); the two are split so callers can pick
either or both. A full staff scraper (departments, phones, etc.) is a
LATER milestone.
"""
await _upsert_one(session, WechatAccount, rec)
async def upsert_staff_from_account(session: AsyncSession, rec: dict[str, Any]) -> None:
"""Best-effort upsert of the Staff row embedded in a wechat-account record.
Maps the account's ``user_id`` + ``user_name`` (plus ``is_on_job`` when
present) onto :class:`Staff` (PK = ``user_id``). Silently no-ops when the
record has no ``user_id`` — there is nothing to key the row on. A
dedicated staff scraper that fills out phone / dept is a LATER milestone.
"""
user_id = rec.get("user_id")
if not user_id:
return
staff_rec: dict[str, Any] = {
"user_id": user_id,
"name": rec.get("user_name"),
"raw_json": rec,
}
if "is_on_job" in rec:
staff_rec["is_on_job"] = rec["is_on_job"]
await _upsert_one(session, Staff, staff_rec)
yunke-server/tests/test_object_store.py
+156
View File
@@ -0,0 +1,156 @@
"""Tests for app.storage.object_store — local-filesystem object store.
Exercises the async hash-dedup store against a real tmp_path on disk. All disk
I/O is dispatched via ``asyncio.to_thread`` so the methods are genuinely async.
"""
from __future__ import annotations
import hashlib
from pathlib import Path
from app.storage.object_store import ObjectStore
# ---------------------------------------------------------------------------
# put_bytes / get_path / exists
# ---------------------------------------------------------------------------
async def test_put_bytes_writes_file_and_exists_reports_true(tmp_path: Path):
"""put_bytes writes the blob; exists() then reports True; path matches key."""
store = ObjectStore(tmp_path)
key = "foo/bar.txt"
data = b"hello"
assert await store.exists(key) is False
path = await store.put_bytes(key, data)
assert path == tmp_path / key
assert path.exists()
assert path.read_bytes() == data
assert await store.exists(key) is True
async def test_get_path_returns_base_plus_key(tmp_path: Path):
"""get_path is a pure path computation (no I/O).
``/`` in the key becomes real path segments (subdirectories).
"""
store = ObjectStore(tmp_path)
assert store.get_path("a/b/c.bin") == tmp_path / "a" / "b" / "c.bin"
assert store.get_path("flat.txt") == tmp_path / "flat.txt"
async def test_put_bytes_overwrites_on_change(tmp_path: Path):
"""A second put_bytes with new content overwrites the file."""
store = ObjectStore(tmp_path)
await store.put_bytes("k", b"v1")
await store.put_bytes("k", b"v2")
assert (tmp_path / "k").read_bytes() == b"v2"
# ---------------------------------------------------------------------------
# put_by_hash (dedup)
# ---------------------------------------------------------------------------
async def test_put_by_hash_returns_sha256_and_path_with_layout(tmp_path: Path):
"""Returns (sha256_hex, base/<2 chars>/<hash><ext>); file content correct."""
store = ObjectStore(tmp_path)
data = b"abc"
expected_hash = hashlib.sha256(data).hexdigest()
h, path = await store.put_by_hash(data, ext=".txt")
assert h == expected_hash
assert path == tmp_path / expected_hash[:2] / (expected_hash + ".txt")
assert path.exists()
assert path.read_bytes() == data
async def test_put_by_hash_dedups_same_bytes_returns_same_path(tmp_path: Path):
"""put_by_hash with identical bytes twice -> same (hash, path)."""
store = ObjectStore(tmp_path)
data = b"identical"
h1, p1 = await store.put_by_hash(data)
h2, p2 = await store.put_by_hash(data)
assert h1 == h2
assert p1 == p2
async def test_put_by_hash_does_not_rewrite_existing_file(tmp_path: Path):
"""Dedup: a second call for an existing hash must NOT touch the file.
Proves the implementation checks ``exists`` before writing. We corrupt the
file between calls; a dedup-aware store must leave the corruption in place
(no rewrite), since the blob is already present.
"""
store = ObjectStore(tmp_path)
data = b"original"
_h1, p1 = await store.put_by_hash(data, ext=".bin")
assert p1.read_bytes() == data
# Corrupt the on-disk file.
p1.write_bytes(b"CORRUPTED")
_h2, p2 = await store.put_by_hash(data, ext=".bin")
assert p2 == p1
# Dedup short-circuit: file content was NOT overwritten.
assert p2.read_bytes() == b"CORRUPTED"
async def test_put_by_hash_distinct_bytes_get_distinct_paths(tmp_path: Path):
"""Different bytes -> different hash -> different path."""
store = ObjectStore(tmp_path)
h1, p1 = await store.put_by_hash(b"aaa")
h2, p2 = await store.put_by_hash(b"bbb")
assert h1 != h2
assert p1 != p2
async def test_put_by_hash_default_ext_empty(tmp_path: Path):
"""Default ext='' -> no suffix appended to the file."""
store = ObjectStore(tmp_path)
data = b"x"
expected_hash = hashlib.sha256(data).hexdigest()
h, p = await store.put_by_hash(data)
assert h == expected_hash
assert p == tmp_path / expected_hash[:2] / expected_hash
# ---------------------------------------------------------------------------
# put_stream
# ---------------------------------------------------------------------------
async def test_put_stream_writes_concatenated_chunks(tmp_path: Path):
"""put_stream consumes an async generator and writes concatenated bytes."""
async def gen():
yield b"chunk1-"
yield b"chunk2-"
yield b"chunk3"
store = ObjectStore(tmp_path)
path = await store.put_stream("streamed/key", gen())
assert path == tmp_path / "streamed" / "key"
assert path.read_bytes() == b"chunk1-chunk2-chunk3"
# ---------------------------------------------------------------------------
# Default base dir (Settings.media_dir)
# ---------------------------------------------------------------------------
def test_object_store_default_base_uses_settings_media_dir(tmp_path: Path, monkeypatch):
"""Omitting base_dir falls back to Settings.media_dir."""
monkeypatch.setenv("MEDIA_DIR", str(tmp_path))
from app.config import get_settings
get_settings.cache_clear()
store = ObjectStore()
assert store.base_dir == get_settings().media_dir
yunke-server/tests/test_repos.py
+397
View File
@@ -0,0 +1,397 @@
"""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()