264 lines
13 KiB
Python
264 lines
13 KiB
Python
"""SQLite migration and attachment boundaries in isolated application contexts."""
|
|
|
|
import io
|
|
import os
|
|
from pathlib import Path
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from threading import Barrier
|
|
from unittest.mock import patch
|
|
from contextlib import closing
|
|
import sqlite3
|
|
|
|
from flask import Flask
|
|
from werkzeug.datastructures import FileStorage
|
|
|
|
from logistics_chat import db, file_service
|
|
from logistics_chat.config import default_config
|
|
|
|
|
|
class StorageTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.directory = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.directory.cleanup)
|
|
self.root = Path(self.directory.name)
|
|
self.app = Flask(__name__)
|
|
self.app.config.update(default_config())
|
|
self.app.config.update(DATABASE=self.root/'history.db', UPLOAD_DIR=self.root/'uploads',
|
|
BACKUP_DIR=self.root/'backups', LOG_DIR=self.root/'logs')
|
|
self.app.teardown_appcontext(db.close_connections)
|
|
self.context = self.app.app_context()
|
|
self.context.push()
|
|
self.addCleanup(self.context.pop)
|
|
db.init_db()
|
|
|
|
def test_concurrent_startup_migrates_legacy_history_once_without_claiming_it(self):
|
|
connection = db.connect()
|
|
connection.executescript('DROP TABLE chat_history; CREATE TABLE chat_history '
|
|
'(id INTEGER PRIMARY KEY, title TEXT NOT NULL, ctime TEXT NOT NULL);'
|
|
"INSERT INTO chat_history VALUES(1,'旧记录','now');")
|
|
connection.close()
|
|
gate = Barrier(4)
|
|
|
|
def initialize():
|
|
with self.app.app_context():
|
|
gate.wait(timeout=5)
|
|
db.init_db()
|
|
|
|
with ThreadPoolExecutor(max_workers=4) as workers:
|
|
for result in workers.map(lambda _: initialize(), range(4)):
|
|
self.assertIsNone(result)
|
|
connection = db.connect()
|
|
self.assertEqual(connection.execute('SELECT title, visitor_id FROM chat_history').fetchall(), [('旧记录', None)])
|
|
self.assertEqual(connection.execute('SELECT COUNT(*) FROM messages WHERE conv_id=1').fetchone()[0], 1)
|
|
connection.close()
|
|
|
|
def test_old_schema_is_migrated_and_legacy_title_backfilled(self):
|
|
connection = db.connect()
|
|
connection.executescript('DROP TABLE chat_history; CREATE TABLE chat_history '
|
|
'(id INTEGER PRIMARY KEY, title TEXT NOT NULL, ctime TEXT NOT NULL);')
|
|
connection.execute('INSERT INTO chat_history VALUES(1, ?, ?)', ('旧会话', '2026-10-20'))
|
|
connection.commit()
|
|
connection.close()
|
|
db.init_db()
|
|
connection = db.connect()
|
|
self.assertEqual(connection.execute('SELECT utime FROM chat_history').fetchone()[0], '2026-10-20')
|
|
self.assertEqual(connection.execute('SELECT role, content FROM messages').fetchone(), ('user', '旧会话'))
|
|
connection.close()
|
|
|
|
def test_business_tables_start_empty_and_reinitialization_preserves_records(self):
|
|
connection = db.connect()
|
|
for table in ('agents', 'orders', 'tickets'):
|
|
self.assertEqual(connection.execute(f'SELECT COUNT(*) FROM {table}').fetchone()[0], 0)
|
|
connection.executescript(
|
|
"INSERT INTO agents VALUES('fixture-agent', 'fixture-password', 'Fixture Agent');"
|
|
"INSERT INTO orders(phone, track_no, status, created) VALUES('123', 'fixture-order', 'active', 'now');"
|
|
"INSERT INTO tickets(content, phone, ctime, status) VALUES('fixture-ticket', '123', 'now', 'pending');"
|
|
)
|
|
connection.commit()
|
|
connection.close()
|
|
db.init_db()
|
|
connection = db.connect()
|
|
self.assertEqual(connection.execute('SELECT username, password, name FROM agents').fetchall(),
|
|
[('fixture-agent', 'fixture-password', 'Fixture Agent')])
|
|
self.assertEqual(connection.execute('SELECT track_no FROM orders').fetchall(), [('fixture-order',)])
|
|
self.assertEqual(connection.execute('SELECT content FROM tickets').fetchall(), [('fixture-ticket',)])
|
|
connection.close()
|
|
|
|
def test_user_retry_is_idempotent_and_card_context_is_readable(self):
|
|
connection = db.connect()
|
|
connection.execute('INSERT INTO chat_history(title, ctime, utime) VALUES(?,?,?)', ('新对话', 'now', 'now'))
|
|
connection.commit()
|
|
connection.close()
|
|
db.save_message(1, 'user', '海运多少钱')
|
|
self.assertEqual(db.save_message(1, 'user', '海运多少钱'), 'duplicate')
|
|
db.save_message(1, 'assistant', '__CARD__{"no":"15252125","status":"派送中"}')
|
|
context = db.load_context(1, '然后呢')
|
|
self.assertEqual(len(context), 3)
|
|
self.assertIn('派送中', context[1]['content'])
|
|
self.assertNotIn('__CARD__', context[1]['content'])
|
|
self.assertEqual(db.last_user_question(1), '海运多少钱')
|
|
|
|
def test_backup_and_retention_only_touch_configured_directories(self):
|
|
backups = self.root/'backups'
|
|
backups.mkdir()
|
|
(backups/'history_20000101.db').write_bytes(b'old')
|
|
db.backup_db()
|
|
self.assertFalse((backups/'history_20000101.db').exists())
|
|
self.assertEqual(len(list(backups.glob('history_*.db'))), 1)
|
|
uploads = self.root/'uploads'
|
|
uploads.mkdir()
|
|
old = uploads/'old.txt'
|
|
old.write_text('old')
|
|
(uploads/'new.txt').write_text('new')
|
|
os.utime(old, (time.time()-8*86400, time.time()-8*86400))
|
|
file_service.cleanup_uploads()
|
|
self.assertFalse(old.exists())
|
|
self.assertEqual((uploads/'new.txt').read_text(), 'new')
|
|
|
|
def test_schema_version_and_pre_migration_snapshot_are_recorded_before_upgrade(self):
|
|
connection = db.connect()
|
|
connection.execute('DROP TABLE IF EXISTS schema_meta')
|
|
connection.commit()
|
|
connection.close()
|
|
db.init_db()
|
|
connection = db.connect()
|
|
self.assertEqual(connection.execute("SELECT value FROM schema_meta WHERE key='schema_version'").fetchone()[0], str(db.SCHEMA_VERSION))
|
|
connection.close()
|
|
snapshots = list((self.root/'backups').glob('pre-migration-*.db'))
|
|
self.assertEqual(len(snapshots), 1)
|
|
self.assertEqual(db.check_database(snapshots[0]), 'ok')
|
|
with closing(sqlite3.connect(snapshots[0])) as snapshot:
|
|
self.assertIsNone(snapshot.execute("SELECT 1 FROM sqlite_master WHERE name='schema_meta'").fetchone())
|
|
|
|
def test_startup_migration_snapshots_keep_only_the_latest_five(self):
|
|
for _ in range(7):
|
|
db.init_db()
|
|
snapshots = sorted((self.root/'backups').glob('pre-migration-*.db'))
|
|
self.assertEqual(len(snapshots), 5)
|
|
self.assertTrue(all(path.is_file() for path in snapshots))
|
|
self.assertTrue(all(db.check_database(path) == 'ok' for path in snapshots))
|
|
|
|
def test_backup_failure_prevents_any_schema_changes(self):
|
|
with closing(db.connect()) as connection:
|
|
connection.execute('DROP TABLE schema_meta')
|
|
connection.commit()
|
|
with patch('logistics_chat.db.snapshot_database', side_effect=OSError('fixture disk full')):
|
|
with self.assertRaisesRegex(OSError, 'disk full'):
|
|
db.init_db()
|
|
with closing(db.connect()) as connection:
|
|
self.assertIsNone(connection.execute("SELECT 1 FROM sqlite_master WHERE name='schema_meta'").fetchone())
|
|
|
|
def test_failed_startup_snapshot_does_not_prune_previous_backups(self):
|
|
backups = self.root/'backups'
|
|
backups.mkdir()
|
|
previous = []
|
|
for index in range(6):
|
|
path = backups/f'pre-migration-20260101-00000{index}-000000-{index:08x}.db'
|
|
db.snapshot_database(self.root/'history.db', path)
|
|
previous.append(path)
|
|
with patch('logistics_chat.db.snapshot_database', side_effect=OSError('fixture disk full')):
|
|
with self.assertRaises(OSError):
|
|
db.init_db()
|
|
self.assertTrue(all(path.exists() for path in previous))
|
|
|
|
def test_failed_migration_rolls_back_schema_and_data_and_keeps_preimage(self):
|
|
with closing(db.connect()) as connection:
|
|
connection.execute('DROP TABLE schema_meta')
|
|
connection.commit()
|
|
def broken(connection):
|
|
connection.execute('CREATE TABLE partial_upgrade(value TEXT)')
|
|
connection.execute("INSERT INTO tickets(content) VALUES('must roll back')")
|
|
raise RuntimeError('fixture migration failed')
|
|
with patch('logistics_chat.db.MIGRATIONS', ((1, broken), *db.MIGRATIONS[1:])):
|
|
with self.assertRaisesRegex(RuntimeError, 'migration failed'):
|
|
db.init_db()
|
|
with closing(db.connect()) as connection:
|
|
self.assertIsNone(connection.execute("SELECT 1 FROM sqlite_master WHERE name='partial_upgrade'").fetchone())
|
|
self.assertEqual(connection.execute('SELECT COUNT(*) FROM tickets').fetchone()[0], 0)
|
|
self.assertEqual(len(list((self.root/'backups').glob('pre-migration-*.db'))), 1)
|
|
|
|
def test_newer_database_version_is_never_downgraded(self):
|
|
with closing(db.connect()) as connection:
|
|
connection.execute("UPDATE schema_meta SET value='999' WHERE key='schema_version'")
|
|
connection.commit()
|
|
with self.assertRaisesRegex(RuntimeError, 'newer'):
|
|
db.init_db()
|
|
with closing(db.connect()) as connection:
|
|
self.assertEqual(connection.execute("SELECT value FROM schema_meta WHERE key='schema_version'").fetchone()[0], '999')
|
|
|
|
def test_missing_migration_step_cannot_mark_database_as_upgraded(self):
|
|
with patch('logistics_chat.db.SCHEMA_VERSION', db.SCHEMA_VERSION + 1):
|
|
with self.assertRaisesRegex(RuntimeError, 'migration'):
|
|
db.init_db()
|
|
with closing(db.connect()) as connection:
|
|
self.assertEqual(connection.execute("SELECT value FROM schema_meta WHERE key='schema_version'").fetchone()[0], str(db.SCHEMA_VERSION))
|
|
|
|
def test_snapshot_backup_is_verified_and_restore_source_is_not_modified(self):
|
|
destination = self.root/'backups'/'manual.db'
|
|
db.snapshot_database(self.root/'history.db', destination)
|
|
self.assertEqual(db.check_database(destination), 'ok')
|
|
original = destination.read_bytes()
|
|
connection = db.connect()
|
|
connection.execute("INSERT INTO chat_history(title,ctime) VALUES('after-backup','now')")
|
|
connection.commit()
|
|
connection.close()
|
|
self.assertEqual(destination.read_bytes(), original)
|
|
|
|
def test_upload_uses_generated_name_and_configured_directory(self):
|
|
upload = FileStorage(stream=io.BytesIO(b'hello'), filename='../report.txt')
|
|
result = file_service.upload_file({'file': upload})
|
|
name = result['url'].rsplit('/', 1)[1]
|
|
self.assertEqual((self.root/'uploads'/name).read_bytes(), b'hello')
|
|
self.assertNotIn('..', name)
|
|
self.assertFalse(result['is_image'])
|
|
|
|
def test_old_attachments_referenced_by_ai_or_human_history_survive_cleanup(self):
|
|
uploads=self.root/'uploads'; uploads.mkdir()
|
|
for name in ('ai.png','human.pdf','unused.txt'):
|
|
file=uploads/name;file.write_bytes(b'fixture')
|
|
os.utime(file,(time.time()-20*86400,time.time()-20*86400))
|
|
connection=db.connect()
|
|
connection.execute("INSERT INTO messages(conv_id,role,content,ctime) VALUES(1,'user','[图片]/uploads/ai.png','now')")
|
|
connection.execute("INSERT INTO human_messages(session_id,side,content,ctime) VALUES(1,'customer','[文件]/uploads/human.pdf','now')")
|
|
connection.commit();connection.close()
|
|
file_service.cleanup_uploads()
|
|
self.assertTrue((uploads/'ai.png').exists())
|
|
self.assertTrue((uploads/'human.pdf').exists())
|
|
self.assertFalse((uploads/'unused.txt').exists())
|
|
|
|
def test_invalid_and_oversized_uploads_are_rejected(self):
|
|
self.assertEqual(file_service.upload_file({})[1], 400)
|
|
denied = FileStorage(stream=io.BytesIO(b'bad'), filename='test.exe')
|
|
self.assertEqual(file_service.upload_file({'file': denied})[1], 400)
|
|
self.app.config['MAX_UPLOAD_MB'] = 1
|
|
oversized = FileStorage(stream=io.BytesIO(b'x'*(1024*1024+1)), filename='test.txt')
|
|
self.assertEqual(file_service.upload_file({'file': oversized})[1], 400)
|
|
self.assertFalse((self.root/'uploads').exists())
|
|
|
|
def test_teardown_rolls_back_uncommitted_work(self):
|
|
with self.app.app_context():
|
|
connection = db.connect()
|
|
connection.execute('INSERT INTO chat_history(title, ctime) VALUES(?,?)', ('uncommitted', 'now'))
|
|
connection = db.connect()
|
|
self.assertEqual(connection.execute('SELECT COUNT(*) FROM chat_history').fetchone()[0], 0)
|
|
connection.close()
|
|
|
|
def test_late_stream_writes_cannot_recreate_messages_for_deleted_conversations(self):
|
|
connection = db.connect()
|
|
cid = connection.execute("INSERT INTO chat_history(title,ctime) VALUES('temporary','now')").lastrowid
|
|
connection.commit()
|
|
db.save_message(cid,'user','发送中的问题')
|
|
connection.execute('DELETE FROM messages WHERE conv_id=?',(cid,))
|
|
connection.execute('DELETE FROM chat_history WHERE id=?',(cid,))
|
|
connection.commit()
|
|
db.save_message(cid,'assistant','迟到的模型回复')
|
|
self.assertEqual(connection.execute('SELECT COUNT(*) FROM messages WHERE conv_id=?',(cid,)).fetchone()[0],0)
|
|
connection.close()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|