test_s6_edit_regenerate.py
7.78 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
"""
S6 消息编辑 + 重新生成 — 测试套件
RED: 这些测试当前应该全部失败（函数和 API 不存在）
"""
import pytest
import subprocess
import sys
import os
import json
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import local_db
BASE_URL = "http://127.0.0.1:5001"
# ── DB 单元测试（使用真实 local_db，操作真实 DB）──
def test_msg_get_exists():
"""TC1: msg_get 返回正确消息"""
# Create a test message via SQL directly
import sqlite3
conn = sqlite3.connect(os.path.expanduser("~/wechat-viewer/circle.db"))
conn.row_factory = sqlite3.Row
c = conn.execute("INSERT INTO conversations (title) VALUES ('__s6_test__')")
cid = c.lastrowid
m = conn.execute("INSERT INTO messages (conversation_id, role, content) VALUES (?, 'user', 'hello s6')", (cid,))
mid = m.lastrowid
conn.commit()
conn.close()
msg = local_db.msg_get(mid)
assert msg is not None, "msg_get returned None"
assert msg["id"] == mid
assert msg["content"] == "hello s6"
# Cleanup
conn = sqlite3.connect(os.path.expanduser("~/wechat-viewer/circle.db"))
conn.execute("DELETE FROM messages WHERE id = ?", (mid,))
conn.execute("DELETE FROM conversations WHERE id = ?", (cid,))
conn.commit()
conn.close()
def test_msg_get_not_exists():
"""TC1b: msg_get 不存在返回 None"""
msg = local_db.msg_get(9999999)
assert msg is None, "Should return None"
def test_msg_update():
"""TC2: msg_update 更新后内容正确"""
import sqlite3
conn = sqlite3.connect(os.path.expanduser("~/wechat-viewer/circle.db"))
conn.row_factory = sqlite3.Row
c = conn.execute("INSERT INTO conversations (title) VALUES ('__s6_test2__')")
cid = c.lastrowid
m = conn.execute("INSERT INTO messages (conversation_id, role, content) VALUES (?, 'user', 'old')", (cid,))
mid = m.lastrowid
conn.commit()
conn.close()
local_db.msg_update(mid, "new content")
conn = sqlite3.connect(os.path.expanduser("~/wechat-viewer/circle.db"))
conn.row_factory = sqlite3.Row
row = conn.execute("SELECT content FROM messages WHERE id = ?", (mid,)).fetchone()
assert row["content"] == "new content"
conn.close()
# Cleanup
conn = sqlite3.connect(os.path.expanduser("~/wechat-viewer/circle.db"))
conn.execute("DELETE FROM messages WHERE id = ?", (mid,))
conn.execute("DELETE FROM conversations WHERE id = ?", (cid,))
conn.commit()
conn.close()
def test_msg_delete_since():
"""TC3: msg_delete_since 删除之后的消息"""
import sqlite3
conn = sqlite3.connect(os.path.expanduser("~/wechat-viewer/circle.db"))
conn.row_factory = sqlite3.Row
c = conn.execute("INSERT INTO conversations (title) VALUES ('__s6_test3__')")
cid = c.lastrowid
m1 = conn.execute("INSERT INTO messages (conversation_id, role, content) VALUES (?, 'user', 'm1')", (cid,))
m2 = conn.execute("INSERT INTO messages (conversation_id, role, content) VALUES (?, 'assistant', 'm2')", (cid,))
m3 = conn.execute("INSERT INTO messages (conversation_id, role, content) VALUES (?, 'user', 'm3')", (cid,))
conn.commit()
conn.close()
local_db.msg_delete_since(cid, m2.lastrowid)
conn = sqlite3.connect(os.path.expanduser("~/wechat-viewer/circle.db"))
conn.row_factory = sqlite3.Row
rows = conn.execute("SELECT id FROM messages WHERE conversation_id = ?", (cid,)).fetchall()
ids = [r["id"] for r in rows]
conn.close()
assert m1.lastrowid in ids
assert m2.lastrowid in ids
assert m3.lastrowid not in ids, f"m3 ({m3.lastrowid}) should be deleted"
# Cleanup
conn = sqlite3.connect(os.path.expanduser("~/wechat-viewer/circle.db"))
conn.execute("DELETE FROM messages WHERE conversation_id = ?", (cid,))
conn.execute("DELETE FROM conversations WHERE id = ?", (cid,))
conn.commit()
conn.close()
# ── API 集成测试（需要 Flask）──
def _curl(method, path, data=None):
cmd = ["curl", "-s", "-X", method, f"{BASE_URL}{path}",
"--connect-timeout", "5", "--max-time", "15"]
if data:
cmd += ["-H", "Content-Type: application/json", "-d", json.dumps(data)]
r = subprocess.run(cmd, capture_output=True, text=True, timeout=20)
try:
return json.loads(r.stdout)
except:
return {"raw": r.stdout[:200]}
def test_api_msg_edit():
"""TC4: PUT edit user message returns ok"""
r1 = _curl("POST", "/api/conversations", {"title": "__s6_edit_test__"})
if not r1.get("ok"):
pytest.skip("Could not create conversation")
cid = r1["data"]["id"]
r2 = _curl("POST", f"/api/conversations/{cid}/messages", {"content": "original"})
if not r2.get("ok"):
pytest.skip("Could not add message")
uid = r2["data"]["user_message"]["id"]
r3 = _curl("PUT", f"/api/conversations/{cid}/messages/{uid}", {"content": "edited s6"})
assert r3.get("ok"), f"Edit failed: {r3}"
r4 = _curl("GET", f"/api/conversations/{cid}/messages")
msgs = r4.get("data", [])
user_msgs = [m for m in msgs if m["id"] == uid]
assert len(user_msgs) == 1
assert user_msgs[0]["content"] == "edited s6"
def test_api_msg_edit_not_user():
"""TC4b: Edit non-user message returns 400"""
r1 = _curl("POST", "/api/conversations", {"title": "__s6_nonuser__"})
cid = r1["data"]["id"]
r2 = _curl("POST", f"/api/conversations/{cid}/messages", {"content": "hi"})
aid = r2["data"]["assistant_message"]["id"]
cmd = ["curl", "-s", "-o", "/dev/stdout", "-w", "\n%{http_code}",
"-X", "PUT", f"{BASE_URL}/api/conversations/{cid}/messages/{aid}",
"-H", "Content-Type: application/json", "-d", '{"content":"hacked"}',
"--connect-timeout", "5", "--max-time", "15"]
r = subprocess.run(cmd, capture_output=True, text=True, timeout=20)
lines = r.stdout.strip().split("\n")
status = lines[-1].strip()
assert status in ("400", "403"), f"Expected 400/403, got {status}"
def test_api_regenerate():
"""TC5: POST regenerate returns SSE stream"""
r1 = _curl("POST", "/api/conversations", {"title": "__s6_regen__"})
cid = r1["data"]["id"]
_curl("POST", f"/api/conversations/{cid}/messages", {"content": "hello"})
cmd = ["curl", "-s", "-o", "/dev/stdout", "-w", "\n%{http_code}",
"-X", "POST", f"{BASE_URL}/api/conversations/{cid}/regenerate",
"--connect-timeout", "5", "--max-time", "60"]
r = subprocess.run(cmd, capture_output=True, text=True, timeout=70)
lines = r.stdout.strip().split("\n")
status = lines[-1].strip()
body = "\n".join(lines[:-1])
assert status == "200", f"Regenerate returned {status}: {body[:200]}"
assert "[DONE]" in body, "SSE should contain [DONE]"
# ── 前端测试 ──
def test_chat_page_has_msgid():
"""TC6: /chat has data-msgid attributes"""
r = subprocess.run(
["curl", "-s", f"{BASE_URL}/chat", "--connect-timeout", "5", "--max-time", "10"],
capture_output=True, text=True, timeout=15
)
assert 'data-msgid' in r.stdout, "data-msgid missing"
def test_regenerate_in_ui():
"""TC7: regenerate function/button exists"""
r = subprocess.run(
["curl", "-s", f"{BASE_URL}/chat", "--connect-timeout", "5", "--max-time", "10"],
capture_output=True, text=True, timeout=15
)
assert 'regenerate' in r.stdout.lower(), "regenerate not in UI"
def test_no_js_escape_bugs():
"""TC8: No JS escape bugs"""
r = subprocess.run(
["curl", "-s", f"{BASE_URL}/chat", "--connect-timeout", "5", "--max-time", "10"],
capture_output=True, text=True, timeout=15
)
html = r.stdout
assert "style=\\\\\\\\'" not in html, "STYLE BACKSLASH BUG"
assert 'onclick=\\\"' not in html, "INLINE ONCLICK ESCAPE BUG"
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "--tb=short"]))