import os import sys import json import sqlite3 import unittest import importlib from pathlib import Path from tempfile import TemporaryDirectory ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT / "src")) class TestChatEditStore(unittest.TestCase): def setUp(self): self._prev_data_dir = os.environ.get("WECHAT_TOOL_DATA_DIR") self._td = TemporaryDirectory() os.environ["WECHAT_TOOL_DATA_DIR"] = self._td.name import wechat_decrypt_tool.app_paths as app_paths import wechat_decrypt_tool.chat_edit_store as chat_edit_store importlib.reload(app_paths) importlib.reload(chat_edit_store) self.app_paths = app_paths self.store = chat_edit_store def tearDown(self): if self._prev_data_dir is None: os.environ.pop("WECHAT_TOOL_DATA_DIR", None) else: os.environ["WECHAT_TOOL_DATA_DIR"] = self._prev_data_dir self._td.cleanup() def test_ensure_schema_creates_db(self): self.store.ensure_schema() db_path = self.app_paths.get_output_dir() / "message_edits.db" self.assertTrue(db_path.exists()) conn = sqlite3.connect(str(db_path)) try: row = conn.execute( "SELECT name FROM sqlite_master WHERE type='table' AND name='message_edits' LIMIT 1" ).fetchone() self.assertIsNotNone(row) finally: conn.close() def test_blob_hex_roundtrip(self): payload = {"a": b"\x00\xff", "nested": {"b": memoryview(b"\x01\x02")}} dumped = self.store.dumps_json_with_blobs(payload) self.assertIn("0x00ff", dumped.lower()) self.assertIn("0x0102", dumped.lower()) loaded = self.store.loads_json_with_blobs(dumped) self.assertEqual(loaded["a"], b"\x00\xff") self.assertEqual(loaded["nested"]["b"], b"\x01\x02") def test_message_id_format_parse(self): mid = self.store.format_message_id("message_0", "Msg_foo", 123) self.assertEqual(mid, "message_0:Msg_foo:123") db, table, local_id = self.store.parse_message_id(mid) self.assertEqual(db, "message_0") self.assertEqual(table, "Msg_foo") self.assertEqual(local_id, 123) with self.assertRaises(ValueError): self.store.parse_message_id("bad") def test_upsert_original_once_does_not_overwrite_snapshot(self): now1 = 1000 now2 = 2000 self.store.upsert_original_once( account="wxid_me", session_id="wxid_you", db="message_0", table_name="Msg_foo", local_id=1, original_msg={"local_id": 1, "message_content": "hello", "compress_content": b"\x01"}, original_resource={"message_id": 9, "packed_info": b"\x02"}, now_ms=now1, ) self.store.upsert_original_once( account="wxid_me", session_id="wxid_you", db="message_0", table_name="Msg_foo", local_id=1, original_msg={"local_id": 1, "message_content": "SHOULD_NOT_OVERWRITE", "compress_content": b"\x03"}, original_resource={"message_id": 9, "packed_info": b"\x04"}, now_ms=now2, ) mid = self.store.format_message_id("message_0", "Msg_foo", 1) item = self.store.get_message_edit("wxid_me", "wxid_you", mid) self.assertIsNotNone(item) self.assertEqual(int(item["first_edited_at"]), now1) self.assertEqual(int(item["last_edited_at"]), now2) self.assertEqual(int(item["edit_count"]), 2) original_msg = self.store.loads_json_with_blobs(item["original_msg_json"]) self.assertEqual(original_msg["message_content"], "hello") self.assertEqual(original_msg["compress_content"], b"\x01") original_res = self.store.loads_json_with_blobs(item["original_resource_json"]) self.assertEqual(int(original_res["message_id"]), 9) self.assertEqual(original_res["packed_info"], b"\x02") def test_update_message_edit_local_id_moves_primary_key(self): self.store.upsert_original_once( account="wxid_me", session_id="wxid_you", db="message_0", table_name="Msg_foo", local_id=10, original_msg={"local_id": 10, "message_content": "hello"}, original_resource=None, now_ms=1234, ) ok = self.store.update_message_edit_local_id( account="wxid_me", session_id="wxid_you", db="message_0", table_name="Msg_foo", old_local_id=10, new_local_id=11, ) self.assertTrue(ok) old_mid = self.store.format_message_id("message_0", "Msg_foo", 10) new_mid = self.store.format_message_id("message_0", "Msg_foo", 11) self.assertIsNone(self.store.get_message_edit("wxid_me", "wxid_you", old_mid)) self.assertIsNotNone(self.store.get_message_edit("wxid_me", "wxid_you", new_mid)) def test_list_sessions_counts(self): self.store.upsert_original_once( account="wxid_me", session_id="u1", db="message_0", table_name="Msg_foo", local_id=1, original_msg={"local_id": 1, "message_content": "a"}, original_resource=None, now_ms=100, ) self.store.upsert_original_once( account="wxid_me", session_id="u1", db="message_0", table_name="Msg_foo", local_id=2, original_msg={"local_id": 2, "message_content": "b"}, original_resource=None, now_ms=200, ) self.store.upsert_original_once( account="wxid_me", session_id="u2", db="message_0", table_name="Msg_foo", local_id=3, original_msg={"local_id": 3, "message_content": "c"}, original_resource=None, now_ms=300, ) stats = self.store.list_sessions("wxid_me") by_sid = {s["session_id"]: s for s in stats} self.assertEqual(int(by_sid["u1"]["msg_count"]), 2) self.assertEqual(int(by_sid["u1"]["last_edited_at"]), 200) self.assertEqual(int(by_sid["u2"]["msg_count"]), 1) self.assertEqual(int(by_sid["u2"]["last_edited_at"]), 300) if __name__ == "__main__": unittest.main()