yunke/yunke-server/tests/test_models.py

"""Tests for app.db.models — ORM model definitions round-trip through SQLite.

These tests build an isolated in-memory SQLite async engine (same pattern as
tests/test_db_base.py) and exercise the real ORM models defined in
app.db.models. They never touch the production MySQL engine.
"""

import datetime as dt

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 (
    ChatGroup,
    Customer,
    Department,
    GroupMember,
    Media,
    Message,
    PhoneCall,
    Sms,
    Staff,
    SyncRun,
    SyncState,
    WechatAccount,
    WhatsappAccount,
    YkAccount,
    message_table_name,
)

ALL_MODELS = (
    Department,
    Staff,
    WechatAccount,
    WhatsappAccount,
    Customer,
    ChatGroup,
    GroupMember,
    Message,
    Media,
    PhoneCall,
    Sms,
    SyncState,
    YkAccount,
    SyncRun,
)

EXPECTED_TABLES = (
    "department",
    "staff",
    "wechat_account",
    "whatsapp_account",
    "customer",
    "chat_group",
    "group_member",
    "message",
    "media",
    "phone_call",
    "sms",
    "sync_state",
    "yk_account",
    "sync_run",
)


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


def test_all_models_extend_base():
    """Every model must subclass Base and declare a __tablename__."""
    for model in ALL_MODELS:
        assert issubclass(model, Base), f"{model.__name__} must extend Base"
        assert isinstance(model.__tablename__, str) and model.__tablename__, (
            f"{model.__name__} must declare __tablename__"
        )


def test_table_names_match_expected():
    """The 14 tables must use the documented snake_case names."""
    for model, expected in zip(ALL_MODELS, EXPECTED_TABLES):
        assert model.__tablename__ == expected, (
            f"{model.__name__}.__tablename__ = {model.__tablename__!r}, "
            f"expected {expected!r}"
        )


async def test_all_tables_create_on_sqlite():
    """All 14 tables must be createable on SQLite (types stay portable)."""
    engine = await _make_engine()
    try:
        async with engine.begin() as conn:
            tables = await conn.run_sync(
                lambda sync_conn: sa.inspect(sync_conn).get_table_names()
            )
        for name in EXPECTED_TABLES:
            assert name in tables, f"missing table: {name}"
    finally:
        await engine.dispose()


async def test_customer_round_trip():
    """A Customer row round-trips its key fields through SQLite."""
    engine = await _make_engine()
    try:
        SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
        async with SessionLocal() as session:
            cust = Customer(
                friend_id="wxid_abc123",
                channel="wechat",
                owner_account_id="wxid_staff",
                friend_nickname="Alice",
                phone="13800138000",
                message_count=42,
            )
            session.add(cust)
            await session.commit()

            fetched = (
                await session.execute(
                    sa.select(Customer).where(Customer.friend_id == "wxid_abc123")
                )
            ).scalar_one()
            assert fetched.friend_id == "wxid_abc123"
            assert fetched.channel == "wechat"
            assert fetched.owner_account_id == "wxid_staff"
            assert fetched.friend_nickname == "Alice"
            assert fetched.phone == "13800138000"
            assert fetched.message_count == 42
            assert fetched.fetched_at is not None
    finally:
        await engine.dispose()


async def test_customer_raw_json_round_trip():
    """raw_json must round-trip an arbitrary dict (stored as JSON)."""
    engine = await _make_engine()
    try:
        SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
        async with SessionLocal() as session:
            payload = {
                "friendWxId": "wxid_abc",
                "list": [1, 2, 3],
                "nested": {"k": "v", "flag": True},
            }
            cust = Customer(
                friend_id="wxid_raw",
                channel="wechat",
                owner_account_id="wxid_staff",
                raw_json=payload,
            )
            session.add(cust)
            await session.commit()

            fetched = (
                await session.execute(
                    sa.select(Customer).where(Customer.friend_id == "wxid_raw")
                )
            ).scalar_one()
            assert fetched.raw_json == payload
    finally:
        await engine.dispose()


async def test_message_natural_pk():
    """Message PK is msg_id (the source msgSvrId / keyId)."""
    engine = await _make_engine()
    try:
        SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
        async with SessionLocal() as session:
            msg = Message(
                msg_id="1000069900000000001",
                channel="wechat",
                owner_account_id="wxid_staff",
                talker_id="wxid_abc",
                is_send=False,
                type=1,
                text="hello",
                timestamp=1700000000000,
            )
            session.add(msg)
            await session.commit()

            fetched = (
                await session.execute(
                    sa.select(Message).where(Message.msg_id == "1000069900000000001")
                )
            ).scalar_one()
            assert fetched.msg_id == "1000069900000000001"
            assert fetched.text == "hello"
            assert fetched.timestamp == 1700000000000
            assert fetched.is_send is False
    finally:
        await engine.dispose()


def test_message_table_name_helper_with_ms_timestamp():
    """message_table_name(channel, ms_epoch) -> message_<channel>_<YYYYMM>."""
    ts_ms = int(
        dt.datetime(2023, 11, 15, 12, 0, 0, tzinfo=dt.UTC).timestamp() * 1000
    )
    assert message_table_name("wechat", ts_ms) == "message_wechat_202311"


def test_message_table_name_helper_with_datetime():
    """message_table_name also accepts a datetime instance."""
    assert (
        message_table_name("whatsapp", dt.datetime(2024, 1, 5, tzinfo=dt.UTC))
        == "message_whatsapp_202401"
    )


async def test_message_composite_index_present():
    """The (channel, owner_account_id, talker_id, timestamp) index must exist."""
    engine = await _make_engine()
    try:
        async with engine.begin() as conn:
            indexes = await conn.run_sync(
                lambda sync_conn: sa.inspect(sync_conn).get_indexes("message")
            )
        names = {idx["name"] for idx in indexes}
        # The composite index name is implementation detail; assert one index
        # covers exactly those four columns.
        cols_to_index = {
            "channel",
            "owner_account_id",
            "talker_id",
            "timestamp",
        }
        assert any(
            set(idx["column_names"]) == cols_to_index for idx in indexes
        ), f"no composite index on {cols_to_index}; got: {names}"
    finally:
        await engine.dispose()


async def test_group_member_composite_pk():
    """GroupMember uses a composite PK (channel, room_id, member_id)."""
    engine = await _make_engine()
    try:
        SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
        async with SessionLocal() as session:
            gm = GroupMember(
                channel="wechat",
                room_id="12345@chatroom",
                member_id="wxid_member1",
                nickname="Bob",
            )
            session.add(gm)
            await session.commit()

            fetched = (
                await session.execute(
                    sa.select(GroupMember).where(
                        GroupMember.room_id == "12345@chatroom",
                        GroupMember.member_id == "wxid_member1",
                        GroupMember.channel == "wechat",
                    )
                )
            ).scalar_one()
            assert fetched.nickname == "Bob"
    finally:
        await engine.dispose()