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/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()