WeChatDataAnalysis/tests/test_chat_message_calendar_heatmap.py

import hashlib
import sqlite3
import sys
import unittest
from datetime import datetime
from pathlib import Path
from tempfile import TemporaryDirectory
from unittest.mock import patch


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))


from wechat_decrypt_tool.routers import chat as chat_router


def _msg_table_name(username: str) -> str:
    md5_hex = hashlib.md5(username.encode("utf-8")).hexdigest()
    return f"Msg_{md5_hex}"


def _seed_message_db(path: Path, *, username: str, rows: list[tuple[int, int]]) -> None:
    """rows: [(create_time, sort_seq), ...]"""
    table = _msg_table_name(username)
    conn = sqlite3.connect(str(path))
    try:
        conn.execute(
            f"""
            CREATE TABLE "{table}"(
                local_id INTEGER PRIMARY KEY AUTOINCREMENT,
                create_time INTEGER,
                sort_seq INTEGER
            )
            """
        )
        for create_time, sort_seq in rows:
            conn.execute(
                f'INSERT INTO "{table}"(create_time, sort_seq) VALUES (?, ?)',
                (int(create_time), int(sort_seq)),
            )
        conn.commit()
    finally:
        conn.close()


def _seed_message_db_full(path: Path, *, username: str, rows: list[tuple[int, int, str]]) -> None:
    """rows: [(create_time, sort_seq, text), ...] - minimal schema for /api/chat/messages/around."""

    table = _msg_table_name(username)
    conn = sqlite3.connect(str(path))
    try:
        conn.execute(
            f"""
            CREATE TABLE "{table}"(
                local_id INTEGER PRIMARY KEY AUTOINCREMENT,
                server_id INTEGER,
                local_type INTEGER,
                sort_seq INTEGER,
                real_sender_id INTEGER,
                create_time INTEGER,
                message_content TEXT,
                compress_content BLOB
            )
            """
        )
        for create_time, sort_seq, text in rows:
            conn.execute(
                f'INSERT INTO "{table}"(server_id, local_type, sort_seq, real_sender_id, create_time, message_content, compress_content) '
                "VALUES (?, ?, ?, ?, ?, ?, ?)",
                (0, 1, int(sort_seq), 0, int(create_time), str(text), None),
            )
        conn.commit()
    finally:
        conn.close()


def _seed_contact_db_minimal(path: Path) -> None:
    conn = sqlite3.connect(str(path))
    try:
        conn.execute(
            """
            CREATE TABLE contact (
                username TEXT,
                remark TEXT,
                nick_name TEXT,
                alias TEXT,
                big_head_url TEXT,
                small_head_url TEXT
            )
            """
        )
        conn.execute(
            """
            CREATE TABLE stranger (
                username TEXT,
                remark TEXT,
                nick_name TEXT,
                alias TEXT,
                big_head_url TEXT,
                small_head_url TEXT
            )
            """
        )
        conn.commit()
    finally:
        conn.close()


class TestChatMessageCalendarHeatmap(unittest.TestCase):
    def test_daily_counts_aggregates_per_day_and_respects_month_range(self):
        with TemporaryDirectory() as td:
            account_dir = Path(td) / "acc"
            account_dir.mkdir(parents=True, exist_ok=True)

            username = "wxid_test_user"

            ts_jan31_23 = int(datetime(2026, 1, 31, 23, 0, 0).timestamp())
            ts_feb01_10 = int(datetime(2026, 2, 1, 10, 0, 0).timestamp())
            ts_feb14_12 = int(datetime(2026, 2, 14, 12, 0, 0).timestamp())

            _seed_message_db(
                account_dir / "message.db",
                username=username,
                rows=[
                    (ts_jan31_23, 0),
                    (ts_feb01_10, 5),
                    (ts_feb01_10, 2),
                    (ts_feb14_12, 0),
                ],
            )

            with patch.object(chat_router, "_resolve_account_dir", return_value=account_dir):
                resp = chat_router.get_chat_message_daily_counts(
                    username=username,
                    year=2026,
                    month=2,
                    account="acc",
                )

            self.assertEqual(resp.get("status"), "success")
            self.assertEqual(resp.get("username"), username)
            self.assertEqual(resp.get("year"), 2026)
            self.assertEqual(resp.get("month"), 2)

            counts = resp.get("counts") or {}
            self.assertEqual(counts.get("2026-02-01"), 2)
            self.assertEqual(counts.get("2026-02-14"), 1)
            self.assertIsNone(counts.get("2026-01-31"))

            self.assertEqual(resp.get("total"), 3)
            self.assertEqual(resp.get("max"), 2)

    def test_anchor_day_picks_earliest_by_create_time_then_sort_seq_then_local_id(self):
        with TemporaryDirectory() as td:
            account_dir = Path(td) / "acc"
            account_dir.mkdir(parents=True, exist_ok=True)

            username = "wxid_test_user"

            ts_jan31_23 = int(datetime(2026, 1, 31, 23, 0, 0).timestamp())
            ts_feb01_10 = int(datetime(2026, 2, 1, 10, 0, 0).timestamp())

            _seed_message_db(
                account_dir / "message.db",
                username=username,
                rows=[
                    (ts_jan31_23, 0),  # local_id = 1
                    (ts_feb01_10, 5),  # local_id = 2
                    (ts_feb01_10, 2),  # local_id = 3  <- expected (sort_seq smaller)
                ],
            )

            with patch.object(chat_router, "_resolve_account_dir", return_value=account_dir):
                resp = chat_router.get_chat_message_anchor(
                    username=username,
                    kind="day",
                    account="acc",
                    date="2026-02-01",
                )

            self.assertEqual(resp.get("status"), "success")
            self.assertEqual(resp.get("kind"), "day")
            self.assertEqual(resp.get("date"), "2026-02-01")
            anchor_id = str(resp.get("anchorId") or "")
            self.assertTrue(anchor_id.startswith("message:"), anchor_id)
            self.assertTrue(anchor_id.endswith(":3"), anchor_id)

    def test_anchor_first_picks_global_earliest(self):
        with TemporaryDirectory() as td:
            account_dir = Path(td) / "acc"
            account_dir.mkdir(parents=True, exist_ok=True)

            username = "wxid_test_user"

            ts_jan31_23 = int(datetime(2026, 1, 31, 23, 0, 0).timestamp())
            ts_feb01_10 = int(datetime(2026, 2, 1, 10, 0, 0).timestamp())

            _seed_message_db(
                account_dir / "message.db",
                username=username,
                rows=[
                    (ts_feb01_10, 2),  # local_id = 1
                    (ts_jan31_23, 0),  # local_id = 2, but earlier create_time -> should win even if local_id bigger
                ],
            )

            with patch.object(chat_router, "_resolve_account_dir", return_value=account_dir):
                resp = chat_router.get_chat_message_anchor(
                    username=username,
                    kind="first",
                    account="acc",
                )

            self.assertEqual(resp.get("status"), "success")
            self.assertEqual(resp.get("kind"), "first")
            anchor_id = str(resp.get("anchorId") or "")
            self.assertTrue(anchor_id.startswith("message:"), anchor_id)
            self.assertTrue(anchor_id.endswith(":2"), anchor_id)

    def test_anchor_day_empty_returns_empty_status(self):
        with TemporaryDirectory() as td:
            account_dir = Path(td) / "acc"
            account_dir.mkdir(parents=True, exist_ok=True)

            username = "wxid_test_user"
            ts_feb01_10 = int(datetime(2026, 2, 1, 10, 0, 0).timestamp())

            _seed_message_db(account_dir / "message.db", username=username, rows=[(ts_feb01_10, 0)])

            with patch.object(chat_router, "_resolve_account_dir", return_value=account_dir):
                resp = chat_router.get_chat_message_anchor(
                    username=username,
                    kind="day",
                    account="acc",
                    date="2026-02-02",
                )

            self.assertEqual(resp.get("status"), "empty")
            self.assertEqual(resp.get("anchorId"), "")

    def test_around_can_span_multiple_message_dbs_for_pagination(self):
        from fastapi import FastAPI
        from fastapi.testclient import TestClient

        with TemporaryDirectory() as td:
            account_dir = Path(td) / "acc"
            account_dir.mkdir(parents=True, exist_ok=True)

            username = "wxid_test_user"
            table = _msg_table_name(username)

            # Anchor in message.db, next message in message_1.db
            _seed_message_db_full(
                account_dir / "message.db",
                username=username,
                rows=[(1000, 0, "A")],  # local_id=1
            )
            _seed_message_db_full(
                account_dir / "message_1.db",
                username=username,
                rows=[(2000, 0, "B")],  # local_id=1
            )
            _seed_contact_db_minimal(account_dir / "contact.db")

            app = FastAPI()
            app.include_router(chat_router.router)
            client = TestClient(app)

            with patch.object(chat_router, "_resolve_account_dir", return_value=account_dir):
                resp = client.get(
                    "/api/chat/messages/around",
                    params={
                        "account": "acc",
                        "username": username,
                        "anchor_id": f"message:{table}:1",
                        "before": 0,
                        "after": 10,
                    },
                )

            self.assertEqual(resp.status_code, 200, resp.text)
            data = resp.json()
            self.assertEqual(data.get("status"), "success")
            self.assertEqual(data.get("username"), username)
            self.assertEqual(data.get("anchorId"), f"message:{table}:1")
            self.assertEqual(data.get("anchorIndex"), 0)

            msgs = data.get("messages") or []
            self.assertEqual(len(msgs), 2)
            self.assertEqual(msgs[0].get("id"), f"message:{table}:1")
            self.assertEqual(msgs[1].get("id"), f"message_1:{table}:1")