250 lines
9.5 KiB
Python
250 lines
9.5 KiB
Python
"""cloud/db.py 单测：SQLite 并发 helper。
|
||
|
||
只测 db.py，不改 server.py / host.py / 其它模块。
|
||
重试逻辑的覆盖方式（本文件写清楚，供协调员核对）：
|
||
- 确定性 fake 方式：FakeConn 前两次抛 database is locked、第三次
|
||
成功，断言重试后成功且退避序列为 0.05/0.1/0.2s；另测重试耗尽
|
||
抛原始错误、非锁错误一次都不重试。
|
||
- 真实并发方式：8 线程各写 20 条（server.py 同款 tasks 表），
|
||
断言无错误、最终 160 条、id 无重复。
|
||
两者都测了，前者证重试逻辑本身正确，后者证真实写冲突不再直接炸掉。
|
||
"""
|
||
|
||
import os
|
||
import sqlite3
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
import unittest
|
||
from pathlib import Path
|
||
from unittest import mock
|
||
|
||
try:
|
||
from cloud import db as db_module
|
||
from cloud.db import (
|
||
connect,
|
||
execute_with_retry,
|
||
is_locked_error,
|
||
read_all,
|
||
transaction,
|
||
)
|
||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||
from cloud import db as db_module
|
||
from cloud.db import (
|
||
connect,
|
||
execute_with_retry,
|
||
is_locked_error,
|
||
read_all,
|
||
transaction,
|
||
)
|
||
|
||
# 与 server.py init_db 的 tasks 表定义逐字段一致（db.py 本身不建表，
|
||
# 这里复刻定义只是为了让并发测试写进同款表，证明 schema 兼容）。
|
||
TASKS_DDL = (
|
||
"CREATE TABLE IF NOT EXISTS tasks("
|
||
"id TEXT PRIMARY KEY, "
|
||
"project_id TEXT, "
|
||
"prompt TEXT, "
|
||
"status TEXT, "
|
||
"idempotency_key TEXT, "
|
||
"depends_on TEXT, "
|
||
"priority INTEGER, "
|
||
"created_by TEXT, "
|
||
"created_at REAL, "
|
||
"updated_at REAL, "
|
||
"last_heartbeat REAL, "
|
||
"UNIQUE(project_id, idempotency_key))"
|
||
)
|
||
|
||
|
||
class FakeConn:
|
||
"""受控 fake：前 fail_times 次 execute 抛锁错误，之后返回哨兵。"""
|
||
|
||
def __init__(self, fail_times, error_text="database is locked"):
|
||
self.fail_times = fail_times
|
||
self.error_text = error_text
|
||
self.calls = 0
|
||
|
||
def execute(self, sql, params=()):
|
||
self.calls += 1
|
||
if self.calls <= self.fail_times:
|
||
raise sqlite3.OperationalError(self.error_text)
|
||
return "ok-sentinel"
|
||
|
||
|
||
class DbHelperTestCase(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "concurrency.db")
|
||
conn = connect(self.db_path)
|
||
try:
|
||
conn.execute(TASKS_DDL)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
# -- 连接属性：WAL + busy_timeout + Row ------------------------------
|
||
def test_connect_sets_wal_busy_timeout_row_factory(self):
|
||
conn = connect(self.db_path)
|
||
try:
|
||
self.assertEqual(conn.row_factory, sqlite3.Row)
|
||
journal = conn.execute(
|
||
"PRAGMA journal_mode").fetchone()[0]
|
||
self.assertEqual(journal, "wal")
|
||
busy = conn.execute("PRAGMA busy_timeout").fetchone()[0]
|
||
self.assertEqual(busy, 5000)
|
||
# Row 行为：按列名访问，与 server.py 的 _task_to_dict 兼容。
|
||
conn.execute(
|
||
"INSERT INTO tasks(id, project_id, prompt, status)"
|
||
" VALUES(?, ?, ?, 'queued')",
|
||
("t-row", "p1", "hello"),
|
||
)
|
||
conn.commit()
|
||
row = read_all(
|
||
conn, "SELECT * FROM tasks WHERE id = ?", ("t-row",))[0]
|
||
self.assertEqual(row["id"], "t-row")
|
||
self.assertEqual(row["prompt"], "hello")
|
||
finally:
|
||
conn.close()
|
||
|
||
def test_connect_custom_busy_timeout(self):
|
||
conn = connect(self.db_path, busy_timeout_ms=1234)
|
||
try:
|
||
busy = conn.execute("PRAGMA busy_timeout").fetchone()[0]
|
||
self.assertEqual(busy, 1234)
|
||
finally:
|
||
conn.close()
|
||
|
||
# -- 锁判定 ---------------------------------------------------------
|
||
def test_is_locked_error_only_for_locked_operational_error(self):
|
||
self.assertTrue(is_locked_error(
|
||
sqlite3.OperationalError("database is locked")))
|
||
self.assertTrue(is_locked_error(
|
||
sqlite3.OperationalError("Database Is Locked (retry)")))
|
||
self.assertFalse(is_locked_error(
|
||
sqlite3.OperationalError("no such table: tasks")))
|
||
self.assertFalse(is_locked_error(ValueError("database is locked")))
|
||
|
||
# -- 重试逻辑：fake 受控方式（确定性）--------------------------------
|
||
def test_retry_succeeds_after_two_locked_failures(self):
|
||
fake = FakeConn(fail_times=2)
|
||
with mock.patch.object(db_module.time, "sleep") as sleep_mock:
|
||
result = execute_with_retry(fake, "INSERT ...", ())
|
||
self.assertEqual(result, "ok-sentinel")
|
||
self.assertEqual(fake.calls, 3) # 1 次初始 + 2 次重试
|
||
self.assertEqual(
|
||
[c.args[0] for c in sleep_mock.call_args_list],
|
||
[0.05, 0.1, 0.2][:2])
|
||
|
||
def test_retry_backoff_sequence_is_exponential(self):
|
||
fake = FakeConn(fail_times=3)
|
||
with mock.patch.object(db_module.time, "sleep") as sleep_mock:
|
||
execute_with_retry(fake, "INSERT ...", ())
|
||
self.assertEqual(
|
||
[c.args[0] for c in sleep_mock.call_args_list],
|
||
[0.05, 0.1, 0.2])
|
||
|
||
def test_retry_exhausted_reraises_original_error(self):
|
||
fake = FakeConn(fail_times=99)
|
||
with mock.patch.object(db_module.time, "sleep"):
|
||
with self.assertRaises(sqlite3.OperationalError) as ctx:
|
||
execute_with_retry(fake, "INSERT ...", ())
|
||
self.assertIn("database is locked", str(ctx.exception))
|
||
# 默认 max_retries=3：1 次初始 + 3 次重试 = 4 次尝试，不多不少。
|
||
self.assertEqual(fake.calls, 4)
|
||
|
||
def test_non_locked_error_is_not_retried(self):
|
||
fake = FakeConn(fail_times=99, error_text="no such table: tasks")
|
||
with mock.patch.object(db_module.time, "sleep") as sleep_mock:
|
||
with self.assertRaises(sqlite3.OperationalError):
|
||
execute_with_retry(fake, "SELECT ...", ())
|
||
self.assertEqual(fake.calls, 1)
|
||
sleep_mock.assert_not_called()
|
||
|
||
# -- 事务：成功提交 / 失败回滚 ---------------------------------------
|
||
def test_transaction_commits_and_rolls_back(self):
|
||
conn = connect(self.db_path)
|
||
try:
|
||
with transaction(conn):
|
||
execute_with_retry(
|
||
conn,
|
||
"INSERT INTO tasks(id, project_id, prompt, status)"
|
||
" VALUES(?, ?, ?, 'queued')",
|
||
("t-ok", "p1", "ok"),
|
||
)
|
||
self.assertEqual(
|
||
len(read_all(conn, "SELECT * FROM tasks")), 1)
|
||
with self.assertRaises(sqlite3.IntegrityError):
|
||
with transaction(conn):
|
||
execute_with_retry(
|
||
conn,
|
||
"INSERT INTO tasks(id, project_id, prompt, status)"
|
||
" VALUES(?, ?, ?, 'queued')",
|
||
("t-ok", "p1", "dup"),
|
||
)
|
||
# 冲突行回滚后仍只有原来 1 条，不留半写。
|
||
self.assertEqual(
|
||
len(read_all(conn, "SELECT * FROM tasks")), 1)
|
||
finally:
|
||
conn.close()
|
||
|
||
# -- 真实并发：8 线程 x 20 条 ----------------------------------------
|
||
def test_concurrent_writes_8_threads_x_20_no_loss(self):
|
||
errors = []
|
||
barrier = threading.Barrier(8)
|
||
|
||
def worker(thread_idx):
|
||
try:
|
||
barrier.wait(timeout=10)
|
||
conn = connect(self.db_path)
|
||
try:
|
||
for i in range(20):
|
||
with transaction(conn):
|
||
execute_with_retry(
|
||
conn,
|
||
"INSERT INTO tasks(id, project_id, prompt,"
|
||
" status, idempotency_key, depends_on,"
|
||
" priority, created_by, created_at,"
|
||
" updated_at, last_heartbeat)"
|
||
" VALUES(?, ?, ?, 'queued', ?, '[]',"
|
||
" 0, 'worker', ?, ?, NULL)",
|
||
("t-%d-%d" % (thread_idx, i),
|
||
"p-concurrent",
|
||
"prompt %d-%d" % (thread_idx, i),
|
||
"key-%d-%d" % (thread_idx, i),
|
||
time.time(), time.time()),
|
||
)
|
||
finally:
|
||
conn.close()
|
||
except Exception as exc: # noqa: BLE001 - 汇总后统一断言
|
||
errors.append(exc)
|
||
|
||
threads = [threading.Thread(target=worker, args=(n,))
|
||
for n in range(8)]
|
||
for t in threads:
|
||
t.start()
|
||
for t in threads:
|
||
t.join(timeout=60)
|
||
self.assertFalse([t for t in threads if t.is_alive()],
|
||
"worker 线程未在 60s 内结束")
|
||
self.assertEqual(errors, [], "并发写不应有任何异常: %r" % (errors,))
|
||
conn = connect(self.db_path)
|
||
try:
|
||
rows = read_all(
|
||
conn,
|
||
"SELECT id FROM tasks WHERE project_id = ?",
|
||
("p-concurrent",))
|
||
finally:
|
||
conn.close()
|
||
self.assertEqual(len(rows), 160)
|
||
self.assertEqual(len({r["id"] for r in rows}), 160)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|