288 lines
15 KiB
Python
288 lines
15 KiB
Python
import base64
|
|
import csv
|
|
import io
|
|
import json
|
|
import sqlite3
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
import urllib.error
|
|
import urllib.request
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from pathlib import Path
|
|
|
|
from wechat_tool import store
|
|
from wechat_tool.importer import import_rows, normalize, parse_time, read_rows
|
|
from wechat_tool.analysis import Redactor, extract_events, local_cards, validate_model_cards, prepare_plan
|
|
from wechat_tool.server import Workbench, sqlite_rows, validate_api, call_model
|
|
from scripts.mac_probe import probe
|
|
from scripts.export_sqlite import export
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
DEMO = json.loads((ROOT / 'examples/demo.json').read_text(encoding='utf-8'))
|
|
|
|
|
|
class CoreTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.db = store.connect(self.temp.name)
|
|
|
|
def tearDown(self):
|
|
self.db.close()
|
|
self.temp.cleanup()
|
|
|
|
def demo(self):
|
|
import_rows(self.db, DEMO, 'demo', {})
|
|
return store.messages(self.db, {})
|
|
|
|
def test_timezone_and_epoch(self):
|
|
self.assertEqual(parse_time('2026-10-01 09:00:00'), '2026-10-01T01:00:00Z')
|
|
self.assertEqual(parse_time('2026-10-01T03:00:00-06:00'), '2026-10-01T09:00:00Z')
|
|
self.assertEqual(parse_time('1767225600000'), '2026-01-01T00:00:00Z')
|
|
with self.assertRaises(ValueError):
|
|
parse_time('2026-10-01', '+14:01')
|
|
|
|
def test_import_dedup_and_rejection(self):
|
|
a = import_rows(self.db, DEMO + [{'content': '缺失字段'}], 'demo', {})
|
|
b = import_rows(self.db, DEMO, 'same-other-name', {})
|
|
self.assertEqual((a['added'], a['rejected'], b['duplicates']), (13, 1, 13))
|
|
self.assertEqual(a['errors'][0]['row'], 14)
|
|
|
|
def test_id_scoped_by_conversation_and_multiline(self):
|
|
rows = [dict(DEMO[0]), dict(DEMO[0], conversation_id='other')]
|
|
self.assertEqual(import_rows(self.db, rows, 'x', {})['added'], 2)
|
|
txt = '[2026-10-01 09:00:00] 张三: 第一行\n第二行\n[2026-10-01 09:02:00] 李四: 回复'
|
|
parsed = read_rows(txt.encode(), 'x.txt')
|
|
self.assertEqual(parsed[0]['content'], '第一行\n第二行')
|
|
self.assertEqual(len(parsed), 2)
|
|
|
|
def test_encodings_and_alias_mapping(self):
|
|
content = '时间,发送人,内容\n2026-10-01 09:00:00,张三,怎么操作？'
|
|
self.assertEqual(read_rows(content.encode('gb18030'), 'x.csv')[0]['发送人'], '张三')
|
|
self.assertEqual(read_rows(content.encode('utf-16'), 'x.csv')[0]['发送人'], '张三')
|
|
msg = normalize({'when': '2026-10-01', 'who': 'a', 'body': 'b'}, {'conversation': 'x', 'mapping': {'timestamp': 'when', 'sender': 'who', 'content': 'body'}})
|
|
self.assertEqual(msg['content'], 'b')
|
|
with self.assertRaises(ValueError):
|
|
normalize(dict(DEMO[0], content=b'compressed'), {})
|
|
|
|
def test_filters_all_private_exclude_and_keyword(self):
|
|
self.demo()
|
|
self.assertEqual(len(store.messages(self.db, {'start': '2026-10-08', 'end': '2026-10-08'})), 2)
|
|
self.assertEqual(len(store.messages(self.db, {'kind': 'private'})), 3)
|
|
self.assertEqual(len(store.messages(self.db, {'conversations': ['demo-test@chatroom']})), 7)
|
|
self.assertEqual(len(store.messages(self.db, {'exclude': ['demo-test@chatroom']})), 6)
|
|
self.assertEqual(len(store.messages(self.db, {'search': '测试比例'})), 3)
|
|
self.assertEqual(len(store.messages(self.db, {'search': '%'})), 0)
|
|
|
|
def test_cross_conversation_clustering_and_media(self):
|
|
rows = self.demo()
|
|
events = extract_events(rows)
|
|
cards = local_cards(events)
|
|
repeated = next(c for c in cards if c['event_count'] == 3)
|
|
self.assertEqual(len(repeated['conversations']), 2)
|
|
self.assertIn('repeat', repeated['signals'])
|
|
self.assertTrue(all(len(c['conversations']) == 1 for c in local_cards(events, False)))
|
|
self.assertTrue(any('非文字' in c['boundaries'] for c in cards))
|
|
self.assertFalse(any('会议' in c['title'] for c in cards))
|
|
|
|
def test_redactor_all_sources_and_custom(self):
|
|
rows = self.demo()
|
|
redactor = Redactor(rows, 'secret', ['虚构客户甲'])
|
|
value = redactor.redact('演示李工 demo-li 演示·测试项目群 虚构客户甲 13800138000 test@example.com https://example.com/?key=abc 密码:hello 192.168.1.1')
|
|
for original in ['演示李工', 'demo-li', '演示·测试项目群', '虚构客户甲', '13800138000', 'test@example.com', 'example.com', 'hello', '192.168.1.1']:
|
|
self.assertNotIn(original, value)
|
|
self.assertEqual(redactor.redact('演示李工'), Redactor(rows, 'secret').redact('演示李工'))
|
|
|
|
def test_batches_do_not_truncate(self):
|
|
rows = self.demo()
|
|
events = extract_events(rows)
|
|
plan = prepare_plan(events, Redactor(rows, 'salt'))
|
|
self.assertEqual(sum(len(b['events']) for b in plan), len(events))
|
|
events[0]['messages'][0]['content'] = '长' * 50000
|
|
with self.assertRaises(ValueError):
|
|
prepare_plan(events, Redactor(rows, 'salt'))
|
|
|
|
def test_model_hallucination_and_missing_event_rejected(self):
|
|
events = extract_events(self.demo())
|
|
valid = {'cards': [dict(title='测试', summary='候选', event_ids=[e['id'] for e in events], evidence_ids=[events[0]['messages'][0]['id']], signals=['repeat'], strategy='轻量记录', status='待核实', answer='待验证', boundaries='需确认')]}
|
|
self.assertEqual(len(validate_model_cards(valid, events)), 1)
|
|
valid['cards'][0]['evidence_ids'] = [999999]
|
|
with self.assertRaises(ValueError):
|
|
validate_model_cards(valid, events)
|
|
with self.assertRaises(ValueError):
|
|
validate_model_cards({'cards': []}, events)
|
|
|
|
def test_sqlite_import_and_readonly_export_with_wal(self):
|
|
source = Path(self.temp.name) / 'source.db'
|
|
src = sqlite3.connect(source)
|
|
src.execute('PRAGMA journal_mode=WAL')
|
|
src.execute('CREATE TABLE chat (id TEXT, content TEXT)')
|
|
src.execute('INSERT INTO chat VALUES (?,?)', ('a', 'hello'))
|
|
src.commit()
|
|
out = Path(self.temp.name) / 'readable.jsonl'
|
|
export(source, 'chat', out)
|
|
self.assertEqual(json.loads(out.read_text())['content'], 'hello')
|
|
self.assertEqual(src.execute('SELECT COUNT(*) FROM chat').fetchone()[0], 1)
|
|
snapshot_path = Path(self.temp.name) / 'snapshot.db'
|
|
snapshot = sqlite3.connect(snapshot_path)
|
|
src.backup(snapshot)
|
|
snapshot.close()
|
|
self.assertEqual(sqlite_rows(snapshot_path.read_bytes(), schema_only=True)[0]['table'], 'chat')
|
|
with self.assertRaises(ValueError):
|
|
sqlite_rows(b'encrypted')
|
|
with self.assertRaises(ValueError):
|
|
export(source, 'chat', source)
|
|
src.close()
|
|
|
|
def test_probe_is_header_only_and_does_not_overwrite(self):
|
|
source = Path(self.temp.name) / 'wx.db'
|
|
source.write_bytes(b'not encrypted header' + b'private' * 100)
|
|
original = source.read_bytes()
|
|
report = probe([self.temp.name])
|
|
found = next(d for d in report['databases'] if d['path'] == str(source.absolute()))
|
|
self.assertEqual(found['format'], 'encrypted_or_unknown')
|
|
self.assertNotIn('private', json.dumps(report))
|
|
self.assertEqual(source.read_bytes(), original)
|
|
|
|
def test_api_endpoint_validation(self):
|
|
for url in ['http://remote.example/v1/chat/completions', 'https://user:pass@example.com/chat/completions', 'https://example.com/chat/completions?key=secret']:
|
|
with self.assertRaises(ValueError):
|
|
validate_api({'endpoint': url, 'model': 'x'})
|
|
with self.assertRaises(ValueError) as error:
|
|
call_model({'endpoint': 'http://127.0.0.1/chat/completions'}, 'private-key\nvalue', {})
|
|
self.assertNotIn('private-key', str(error.exception))
|
|
|
|
|
|
class FakeModel(BaseHTTPRequestHandler):
|
|
calls = []
|
|
invalid = False
|
|
|
|
def log_message(self, *args):
|
|
pass
|
|
|
|
def do_POST(self):
|
|
body = json.loads(self.rfile.read(int(self.headers['Content-Length'])))
|
|
type(self).calls.append(body)
|
|
batch = json.loads(body['messages'][1]['content'])
|
|
cards = [dict(title='脱敏主题', summary='提炼的经验线索', event_ids=[e['id']], evidence_ids=[e['messages'][0]['id']], signals=[], status='待核实', strategy='轻量记录', answer='需要人工验证', boundaries='适用条件待确认') for e in batch['events']]
|
|
if self.invalid:
|
|
cards[0]['evidence_ids'] = [987654321]
|
|
value = json.dumps({'choices': [{'message': {'content': json.dumps({'cards': cards}, ensure_ascii=False)}}]}).encode()
|
|
self.send_response(200)
|
|
self.send_header('Content-Type', 'application/json')
|
|
self.send_header('Content-Length', str(len(value)))
|
|
self.end_headers()
|
|
self.wfile.write(value)
|
|
|
|
|
|
class HttpTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.app = Workbench(('127.0.0.1', 0), self.temp.name)
|
|
self.thread = threading.Thread(target=self.app.serve_forever, daemon=True)
|
|
self.thread.start()
|
|
self.model = ThreadingHTTPServer(('127.0.0.1', 0), FakeModel)
|
|
threading.Thread(target=self.model.serve_forever, daemon=True).start()
|
|
FakeModel.calls, FakeModel.invalid = [], False
|
|
self.base = f'http://127.0.0.1:{self.app.server_port}'
|
|
|
|
def tearDown(self):
|
|
self.app.shutdown()
|
|
self.app.server_close()
|
|
self.model.shutdown()
|
|
self.model.server_close()
|
|
self.thread.join()
|
|
self.temp.cleanup()
|
|
|
|
def request(self, path, body=None, headers=None, raw=False):
|
|
opts = {'Content-Type': 'application/json', 'X-Workbench-Token': self.app.token}
|
|
if headers:
|
|
opts.update(headers)
|
|
req = urllib.request.Request(self.base+path, data=json.dumps(body).encode() if body is not None else None, headers=opts)
|
|
with urllib.request.urlopen(req) as resp:
|
|
data = resp.read()
|
|
return data if raw else json.loads(data)
|
|
|
|
def configure(self):
|
|
return self.request('/api/settings', {'api': {'endpoint': f'http://127.0.0.1:{self.model.server_port}/v1/chat/completions', 'model': 'mock'}, 'key': 'test-secret-key', 'custom': ['虚构客户甲'], 'utc_offset': '+08:00'})
|
|
|
|
def wait_run(self, rid):
|
|
for _ in range(60):
|
|
result = self.request('/api/run?id='+rid)
|
|
if result['status'] != 'running':
|
|
return result
|
|
time.sleep(.05)
|
|
self.fail('mock model task timed out')
|
|
|
|
def test_csrf_and_host_protection(self):
|
|
with self.assertRaises(urllib.error.HTTPError) as error:
|
|
self.request('/api/demo', {}, {'X-Workbench-Token': ''})
|
|
self.assertEqual(error.exception.code, 403)
|
|
with self.assertRaises(urllib.error.HTTPError):
|
|
self.request('/api/overview', headers={'Host': 'evil.example'})
|
|
with self.assertRaises(urllib.error.HTTPError):
|
|
self.request('/api/demo', {}, {'Origin': 'https://evil.example'})
|
|
|
|
def test_complete_local_review_split_merge_export(self):
|
|
self.request('/api/demo', {})
|
|
run = self.request('/api/analyze', {'filters': {'start': '2026-10-01', 'end': '2026-10-31'}, 'cross': True})
|
|
result = self.request('/api/run?id='+run['id'])
|
|
card = next(c for c in result['cards'] if c['event_count'] == 3)
|
|
self.request('/api/card', {'run_id': run['id'], 'ids': [card['id']], 'action': 'edit', 'title': '=HYPERLINK("bad")', 'reviewed': True})
|
|
csv_text = self.request('/api/export?id='+run['id']+'&format=csv', raw=True).decode('utf-8-sig')
|
|
self.assertIn("'=HYPERLINK", csv_text)
|
|
self.assertIn('微信知识需求清单', self.request('/api/export?id='+run['id']+'&format=md', raw=True).decode())
|
|
context = self.request('/api/context?run_id='+run['id']+'&card_id='+card['id'])
|
|
self.assertEqual(len(context), 3)
|
|
# Later imports in the same time window must not rewrite historical evidence.
|
|
extra = dict(DEMO[1], id='new-id', timestamp='2026-10-01 09:04:00', content='后来补充的记录')
|
|
self.request('/api/import', {'filename': 'late.json', 'data': base64.b64encode(json.dumps([extra]).encode()).decode()})
|
|
self.assertEqual(context, self.request('/api/context?run_id='+run['id']+'&card_id='+card['id']))
|
|
self.request('/api/card', {'run_id': run['id'], 'ids': [card['id']], 'action': 'split', 'event_ids': [card['event_ids'][0]]})
|
|
split = self.request('/api/run?id='+run['id'])
|
|
self.assertEqual(len(split['cards']), len(result['cards'])+1)
|
|
ids = [c['id'] for c in split['cards'][-2:]]
|
|
self.request('/api/card', {'run_id': run['id'], 'ids': ids, 'action': 'merge'})
|
|
merged = self.request('/api/run?id='+run['id'])
|
|
self.assertEqual(len(merged['cards']), len(result['cards']))
|
|
|
|
def test_preview_exact_send_key_not_persisted(self):
|
|
self.request('/api/demo', {})
|
|
self.configure()
|
|
preview = self.request('/api/preview', {'filters': {}, 'cross': True})
|
|
self.assertFalse(FakeModel.calls)
|
|
serialized = json.dumps(preview, ensure_ascii=False)
|
|
for original in ['演示李工', 'demo-li', '虚构客户甲', '13800138000', 'test@example.com', 'test-secret-key', 'internal.example.com']:
|
|
self.assertNotIn(original, serialized)
|
|
job = self.request('/api/send', {'plan_id': preview['plan_id']})
|
|
result = self.wait_run(job['id'])
|
|
self.assertEqual(result['status'], 'complete', result['error'])
|
|
self.assertEqual(FakeModel.calls, preview['requests'])
|
|
with self.assertRaises(urllib.error.HTTPError):
|
|
self.request('/api/send', {'plan_id': preview['plan_id']})
|
|
db = self.app.db()
|
|
self.assertNotIn('test-secret-key', '\n'.join(db.iterdump()))
|
|
db.close()
|
|
|
|
def test_invalid_model_evidence_fails_atomically(self):
|
|
self.request('/api/demo', {})
|
|
self.configure()
|
|
FakeModel.invalid = True
|
|
preview = self.request('/api/preview', {'filters': {}})
|
|
result = self.wait_run(self.request('/api/send', {'plan_id': preview['plan_id']})['id'])
|
|
self.assertEqual(result['status'], 'failed')
|
|
self.assertEqual(result['cards'], [])
|
|
self.assertIn('证据', result['error'])
|
|
|
|
def test_json_upload_and_rejected_encrypted_db(self):
|
|
payload = {'filename': 'chat.json', 'data': base64.b64encode(json.dumps(DEMO).encode()).decode(), 'options': {}}
|
|
self.assertEqual(self.request('/api/import', payload)['added'], 13)
|
|
payload['filename'] = 'encrypted.db'
|
|
with self.assertRaises(urllib.error.HTTPError):
|
|
self.request('/api/schema', payload)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|