# -*- coding: utf-8 -*- import unittest import sys import os import shutil import sqlite3 import tempfile import threading from datetime import date from unittest.mock import patch import pandas as pd from sqlalchemy import and_, create_engine as sqlalchemy_create_engine, inspect, select from sqlalchemy.sql import func # Ensure src module can be imported sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) from src.config import Config from src.storage import Base, CURRENT_SCHEMA_VERSION, DatabaseManager, DatabaseSchemaMigration, StockDaily from src.services import stock_list_parser as stock_list_parser_module class TestStorage(unittest.TestCase): @staticmethod def _list_sqlite_indexes(db_path: str, table_name: str) -> dict[str, list[str]]: with sqlite3.connect(db_path) as conn: indexes = {} for row in conn.execute(f"PRAGMA index_list({table_name})").fetchall(): index_name = row[1] indexes[index_name] = [ index_info[2] for index_info in conn.execute( f"PRAGMA index_xinfo({index_name})" ).fetchall() if index_info[2] is not None and int(index_info[5]) == 1 ] return indexes @staticmethod def _list_sqlite_unique_indexes(db_path: str, table_name: str) -> dict[str, list[str]]: with sqlite3.connect(db_path) as conn: rows = conn.execute(f"PRAGMA index_list({table_name})").fetchall() unique_indexes = {} for row in rows: if int(row[2]) != 1: continue index_name = row[1] index_columns = [] for index_info in conn.execute(f"PRAGMA index_xinfo({index_name})").fetchall(): column_name = index_info[2] if column_name is not None: index_columns.append(column_name) unique_indexes[index_name] = index_columns return unique_indexes def test_legacy_intelligence_items_url_unique_index_rebuilds_without_collision(self) -> None: temp_dir = tempfile.TemporaryDirectory() db_path = os.path.join(temp_dir.name, "legacy_intel.sqlite") try: with sqlite3.connect(db_path) as conn: conn.execute( """CREATE TABLE intelligence_sources ( id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL UNIQUE, source_type TEXT NOT NULL DEFAULT 'rss', url TEXT NOT NULL, scope_type TEXT NOT NULL DEFAULT 'market', scope_value TEXT, market TEXT NOT NULL DEFAULT 'cn', enabled INTEGER NOT NULL DEFAULT 1, last_status TEXT, last_error TEXT, last_fetched_at DATETIME, created_at DATETIME, updated_at DATETIME )""" ) conn.execute( """CREATE TABLE intelligence_items ( id INTEGER PRIMARY KEY AUTOINCREMENT, source_id INTEGER, source_name TEXT, source_type TEXT NOT NULL DEFAULT 'rss', title TEXT NOT NULL, summary TEXT, url TEXT NOT NULL, source TEXT, published_at DATETIME, fetched_at DATETIME, scope_type TEXT NOT NULL DEFAULT 'market', scope_value TEXT, market TEXT NOT NULL DEFAULT 'cn', raw_payload TEXT )""" ) conn.execute("CREATE UNIQUE INDEX uix_intelligence_item_url_legacy ON intelligence_items(url)") conn.execute("CREATE INDEX ix_intel_item_scope_time ON intelligence_items(scope_type, scope_value, market, published_at)") conn.execute("CREATE INDEX ix_intel_item_fetch_time ON intelligence_items(fetched_at)") conn.execute("INSERT INTO intelligence_sources (name, url) VALUES ('legacy', 'https://legacy.example.com/rss.xml')") source_id = conn.execute("SELECT id FROM intelligence_sources WHERE name='legacy'").fetchone()[0] conn.executemany( """INSERT INTO intelligence_items ( source_id, source_name, source_type, title, summary, url, source, published_at, fetched_at, scope_type, scope_value, market, raw_payload ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", [ (source_id, 'legacy-source', 'rss', 'A', 'legacy-a', 'https://legacy.example.com/a', 'legacy', '2026-01-01 00:00:00', '2026-01-01 00:00:00', 'market', None, 'cn', None), (source_id, 'legacy-source', 'rss', 'B', 'legacy-b', 'https://legacy.example.com/b', 'legacy', '2026-01-02 00:00:00', '2026-01-02 00:00:00', 'market', None, 'cn', None), ], ) unique_indexes_before = self._list_sqlite_unique_indexes(db_path, "intelligence_items") self.assertIn("uix_intelligence_item_url_legacy", unique_indexes_before) self.assertEqual(unique_indexes_before["uix_intelligence_item_url_legacy"], ["url"]) DatabaseManager.reset_instance() Config.reset_instance() DatabaseManager(db_url=f"sqlite:///{db_path}") unique_indexes_after = self._list_sqlite_unique_indexes(db_path, "intelligence_items") self.assertNotIn("uix_intelligence_item_url_legacy", unique_indexes_after) self.assertIn("uix_intel_item_scope", unique_indexes_after) self.assertEqual( unique_indexes_after["uix_intel_item_scope"], ["source_id", "url", "scope_type", "scope_value", "market"], ) with sqlite3.connect(db_path) as conn: table_count = conn.execute("SELECT COUNT(*) FROM intelligence_items").fetchone()[0] temp_tables = conn.execute( "SELECT name FROM sqlite_master WHERE type='table' AND name LIKE 'intelligence_items_recreate_tmp_%'" ).fetchall() self.assertEqual(table_count, 2) self.assertEqual(temp_tables, []) finally: DatabaseManager.reset_instance() Config.reset_instance() temp_dir.cleanup() def test_database_initialization_records_schema_version(self): DatabaseManager.reset_instance() db = DatabaseManager(db_url="sqlite:///:memory:") with db.get_session() as session: row = session.get(DatabaseSchemaMigration, CURRENT_SCHEMA_VERSION) self.assertIsNotNone(row) self.assertEqual(row.version, CURRENT_SCHEMA_VERSION) self.assertIn("metadata.create_all", row.description) DatabaseManager.reset_instance() def test_schema_migration_record_is_idempotent(self): DatabaseManager.reset_instance() db = DatabaseManager(db_url="sqlite:///:memory:") db._ensure_schema_migration_record() db._ensure_schema_migration_record() with db.get_session() as session: count = session.execute( select(func.count()).select_from(DatabaseSchemaMigration) ).scalar_one() self.assertEqual(count, 1) DatabaseManager.reset_instance() def test_fresh_decision_signal_schema_has_profile_indexes(self): DatabaseManager.reset_instance() temp_dir = tempfile.TemporaryDirectory() db_path = os.path.join(temp_dir.name, "fresh_decision_profile.db") try: DatabaseManager(db_url=f"sqlite:///{db_path}") indexes = self._list_sqlite_indexes(db_path, "decision_signals") self.assertEqual( indexes.get("ix_decision_signals_decision_profile"), ["decision_profile"], ) self.assertEqual( indexes.get("ix_decision_signal_market_stock_profile_created"), ["market", "stock_code", "decision_profile", "created_at"], ) self.assertEqual( indexes.get( "ix_decision_signal_report_type_market_stock_profile_action_horizon_phase" ), [ "source_report_id", "source_type", "market", "stock_code", "decision_profile", "action", "horizon", "market_phase", ], ) self.assertEqual( indexes.get( "ix_decision_signal_trace_type_market_stock_profile_action_horizon_phase" ), [ "trace_id", "source_type", "market", "stock_code", "decision_profile", "action", "horizon", "market_phase", ], ) finally: DatabaseManager.reset_instance() Config.reset_instance() temp_dir.cleanup() def test_decision_signal_profile_migration_adds_column_indexes_and_closed_stats(self): DatabaseManager.reset_instance() temp_dir = tempfile.TemporaryDirectory() db_path = os.path.join(temp_dir.name, "legacy_decision_profile.db") deeply_nested_json = "[" * 10_000 + "]" * 10_000 try: with sqlite3.connect(db_path) as conn: conn.execute( """CREATE TABLE decision_signals ( id INTEGER PRIMARY KEY AUTOINCREMENT, stock_code TEXT, market TEXT, source_type TEXT, source_report_id INTEGER, trace_id TEXT, action TEXT, horizon TEXT, market_phase TEXT, created_at DATETIME, metadata_json TEXT )""" ) conn.execute( "CREATE INDEX ix_decision_signal_report_type_market_stock_action_horizon_phase " "ON decision_signals " "(source_report_id, source_type, market, stock_code, action, horizon, market_phase)" ) conn.execute( "CREATE INDEX ix_decision_signal_trace_type_market_stock_action_horizon_phase " "ON decision_signals " "(trace_id, source_type, market, stock_code, action, horizon, market_phase)" ) conn.executemany( """INSERT INTO decision_signals ( stock_code, market, source_type, source_report_id, trace_id, action, horizon, market_phase, created_at, metadata_json ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", [ ("600519", "cn", "analysis", 1, "trace-1", "buy", "3d", "intraday", "2026-01-01", '{"decision_profile":"balanced"}'), ("600519", "cn", "analysis", 2, "trace-2", "buy", "3d", "intraday", "2026-01-01", None), ("600519", "cn", "analysis", 3, "trace-3", "buy", "3d", "intraday", "2026-01-01", ""), ("600519", "cn", "analysis", 4, "trace-4", "buy", "3d", "intraday", "2026-01-01", " "), ("600519", "cn", "analysis", 5, "trace-5", "buy", "3d", "intraday", "2026-01-01", "{not-json"), ("600519", "cn", "analysis", 6, "trace-6", "buy", "3d", "intraday", "2026-01-01", sqlite3.Binary(b"\xff")), ("600519", "cn", "analysis", 7, "trace-7", "buy", "3d", "intraday", "2026-01-01", "null"), ("600519", "cn", "analysis", 8, "trace-8", "buy", "3d", "intraday", "2026-01-01", "[]"), ("600519", "cn", "analysis", 9, "trace-9", "buy", "3d", "intraday", "2026-01-01", '"balanced"'), ("600519", "cn", "analysis", 10, "trace-10", "buy", "3d", "intraday", "2026-01-01", "1"), ("600519", "cn", "analysis", 11, "trace-11", "buy", "3d", "intraday", "2026-01-01", "{}"), ("600519", "cn", "analysis", 12, "trace-12", "buy", "3d", "intraday", "2026-01-01", '{"decision_profile":null}'), ("600519", "cn", "analysis", 13, "trace-13", "buy", "3d", "intraday", "2026-01-01", '{"decision_profile":""}'), ("600519", "cn", "analysis", 14, "trace-14", "buy", "3d", "intraday", "2026-01-01", '{"decision_profile":" "}'), ("600519", "cn", "analysis", 15, "trace-15", "buy", "3d", "intraday", "2026-01-01", '{"decision_profile":"reckless"}'), ("600519", "cn", "analysis", 16, "trace-16", "buy", "3d", "intraday", "2026-01-01", deeply_nested_json), ], ) with self.assertLogs("src.storage", level="INFO") as logs: DatabaseManager(db_url=f"sqlite:///{db_path}") with sqlite3.connect(db_path) as conn: columns = {row[1] for row in conn.execute("PRAGMA table_info(decision_signals)").fetchall()} rows = conn.execute( "SELECT id, decision_profile FROM decision_signals ORDER BY id" ).fetchall() self.assertIn("decision_profile", columns) self.assertEqual(rows[0], (1, "balanced")) self.assertTrue(all(profile is None for _, profile in rows[1:])) indexes = self._list_sqlite_indexes(db_path, "decision_signals") expected_indexes = { "ix_decision_signals_decision_profile": ["decision_profile"], "ix_decision_signal_market_stock_profile_created": [ "market", "stock_code", "decision_profile", "created_at", ], "ix_decision_signal_report_type_market_stock_profile_action_horizon_phase": [ "source_report_id", "source_type", "market", "stock_code", "decision_profile", "action", "horizon", "market_phase", ], "ix_decision_signal_trace_type_market_stock_profile_action_horizon_phase": [ "trace_id", "source_type", "market", "stock_code", "decision_profile", "action", "horizon", "market_phase", ], } for index_name, index_columns in expected_indexes.items(): self.assertEqual(indexes.get(index_name), index_columns) self.assertEqual( indexes.get("ix_decision_signal_report_type_market_stock_action_horizon_phase"), [ "source_report_id", "source_type", "market", "stock_code", "action", "horizon", "market_phase", ], ) self.assertEqual( indexes.get("ix_decision_signal_trace_type_market_stock_action_horizon_phase"), [ "trace_id", "source_type", "market", "stock_code", "action", "horizon", "market_phase", ], ) log_text = "\n".join(logs.output) self.assertIn("candidate_count=16", log_text) self.assertIn("backfilled_count=1", log_text) self.assertIn("guard_skipped_count=0", log_text) self.assertIn("missing_metadata_count=1", log_text) self.assertIn("missing_profile_count=4", log_text) self.assertIn("invalid_json_count=5", log_text) self.assertIn("non_object_count=4", log_text) self.assertIn("invalid_profile_count=1", log_text) self.assertIn("skipped_existing_profile_count=0", log_text) DatabaseManager.reset_instance() with self.assertLogs("src.storage", level="INFO") as second_logs: DatabaseManager(db_url=f"sqlite:///{db_path}") second_log_text = "\n".join(second_logs.output) self.assertIn("candidate_count=15", second_log_text) self.assertIn("backfilled_count=0", second_log_text) self.assertIn("guard_skipped_count=0", second_log_text) self.assertIn("missing_metadata_count=1", second_log_text) self.assertIn("missing_profile_count=4", second_log_text) self.assertIn("invalid_json_count=5", second_log_text) self.assertIn("non_object_count=4", second_log_text) self.assertIn("invalid_profile_count=1", second_log_text) self.assertIn("skipped_existing_profile_count=1", second_log_text) finally: DatabaseManager.reset_instance() Config.reset_instance() temp_dir.cleanup() def test_decision_signal_profile_migration_runs_when_column_already_exists(self): DatabaseManager.reset_instance() temp_dir = tempfile.TemporaryDirectory() db_path = os.path.join(temp_dir.name, "existing_decision_profile.db") try: with sqlite3.connect(db_path) as conn: conn.execute( """CREATE TABLE decision_signals ( id INTEGER PRIMARY KEY AUTOINCREMENT, stock_code TEXT, market TEXT, source_type TEXT, source_report_id INTEGER, trace_id TEXT, action TEXT, horizon TEXT, market_phase TEXT, created_at DATETIME, metadata_json TEXT, decision_profile VARCHAR(16) )""" ) conn.executemany( "INSERT INTO decision_signals (metadata_json, decision_profile) VALUES (?, ?)", [ ('{"decision_profile":"aggressive"}', None), (None, None), ('{"decision_profile":"balanced"}', "conservative"), ('{"decision_profile":"balanced"}', ""), ], ) with self.assertLogs("src.storage", level="INFO") as logs: DatabaseManager(db_url=f"sqlite:///{db_path}") with sqlite3.connect(db_path) as conn: profiles = conn.execute( "SELECT decision_profile FROM decision_signals ORDER BY id" ).fetchall() self.assertEqual( profiles, [("aggressive",), (None,), ("conservative",), ("",)], ) log_text = "\n".join(logs.output) self.assertIn("candidate_count=2", log_text) self.assertIn("backfilled_count=1", log_text) self.assertIn("missing_metadata_count=1", log_text) self.assertIn("skipped_existing_profile_count=2", log_text) finally: DatabaseManager.reset_instance() Config.reset_instance() temp_dir.cleanup() def test_decision_signal_profile_migration_fails_when_column_inspection_fails(self): class BrokenInspector: def has_table(self, _table_name: str) -> bool: return True def get_columns(self, _table_name: str): raise RuntimeError("inspection failed") DatabaseManager.reset_instance() try: with patch("src.storage.inspect", return_value=BrokenInspector()): with self.assertLogs("src.storage", level="ERROR") as logs: with self.assertRaises(RuntimeError): DatabaseManager(db_url="sqlite:///:memory:") self.assertIn( "profile migration cannot continue safely", "\n".join(logs.output), ) finally: DatabaseManager.reset_instance() Config.reset_instance() def test_schema_migration_record_handles_concurrent_initialization(self): DatabaseManager.reset_instance() temp_dir = tempfile.TemporaryDirectory() db_path = os.path.join(temp_dir.name, "schema_migration_race.db") db = DatabaseManager(db_url=f"sqlite:///{db_path}") worker_count = 8 barrier = threading.Barrier(worker_count) errors = [] state_lock = threading.Lock() with db.get_session() as session: session.query(DatabaseSchemaMigration).delete() session.commit() def ensure_record() -> None: try: barrier.wait(timeout=5) db._ensure_schema_migration_record() except Exception as exc: with state_lock: errors.append(exc) threads = [threading.Thread(target=ensure_record) for _ in range(worker_count)] for thread in threads: thread.start() for thread in threads: thread.join(timeout=5) with db.get_session() as session: rows = session.execute(select(DatabaseSchemaMigration)).scalars().all() self.assertFalse(any(thread.is_alive() for thread in threads)) self.assertEqual(errors, []) self.assertEqual(len(rows), 1) self.assertEqual(rows[0].version, CURRENT_SCHEMA_VERSION) DatabaseManager.reset_instance() temp_dir.cleanup() def test_parse_sniper_value(self): """测试解析狙击点位数值""" # 1. 正常数值 self.assertEqual(DatabaseManager._parse_sniper_value(100), 100.0) self.assertEqual(DatabaseManager._parse_sniper_value(100.5), 100.5) self.assertEqual(DatabaseManager._parse_sniper_value("100"), 100.0) self.assertEqual(DatabaseManager._parse_sniper_value("100.5"), 100.5) # 2. 包含中文描述和"元" self.assertEqual(DatabaseManager._parse_sniper_value("建议在 100 元附近买入"), 100.0) self.assertEqual(DatabaseManager._parse_sniper_value("价格：100.5元"), 100.5) # 3. 包含干扰数字（修复的Bug场景） # 之前 "MA5" 会被错误提取为 5.0，现在应该提取 "元" 前面的 100 text_bug = "无法给出。需等待MA5数据恢复，在股价回踩MA5且乖离率<2%时考虑100元" self.assertEqual(DatabaseManager._parse_sniper_value(text_bug), 100.0) # 4. 更多干扰场景 text_complex = "MA10为20.5，建议在30元买入" self.assertEqual(DatabaseManager._parse_sniper_value(text_complex), 30.0) text_multiple = "支撑位10元，阻力位20元" # 应该提取最后一个"元"前面的数字，即20，或者更复杂的逻辑？ # 当前逻辑是找最后一个冒号，然后找之后的第一个"元"，提取中间的数字。 # 测试没有冒号的情况 self.assertEqual(DatabaseManager._parse_sniper_value("30元"), 30.0) # 测试多个数字在"元"之前 self.assertEqual(DatabaseManager._parse_sniper_value("MA5 10 20元"), 20.0) # 5. Fallback: no "元" character — extracts last non-MA number self.assertEqual(DatabaseManager._parse_sniper_value("102.10-103.00（MA5附近）"), 103.0) self.assertEqual(DatabaseManager._parse_sniper_value("97.62-98.50（MA10附近）"), 98.5) self.assertEqual(DatabaseManager._parse_sniper_value("93.40下方（MA20支撑）"), 93.4) self.assertEqual(DatabaseManager._parse_sniper_value("108.00-110.00（前期高点阻力）"), 110.0) # 6. 无效输入 self.assertIsNone(DatabaseManager._parse_sniper_value(None)) self.assertIsNone(DatabaseManager._parse_sniper_value("")) self.assertIsNone(DatabaseManager._parse_sniper_value("没有数字")) self.assertIsNone(DatabaseManager._parse_sniper_value("MA5但没有元")) # 7. 回归：括号内技术指标数字不应被提取 self.assertNotEqual(DatabaseManager._parse_sniper_value("1.52-1.53 (回踩MA5/10附近)"), 10.0) self.assertNotEqual(DatabaseManager._parse_sniper_value("1.55-1.56(MA5/M20支撑)"), 20.0) self.assertNotEqual(DatabaseManager._parse_sniper_value("1.49-1.50(MA60附近企稳)"), 60.0) # 验证正确值在区间内 self.assertIn(DatabaseManager._parse_sniper_value("1.52-1.53 (回踩MA5/10附近)"), [1.52, 1.53]) self.assertIn(DatabaseManager._parse_sniper_value("1.55-1.56(MA5/M20支撑)"), [1.55, 1.56]) self.assertIn(DatabaseManager._parse_sniper_value("1.49-1.50(MA60附近企稳)"), [1.49, 1.50]) def test_get_chat_sessions_prefix_is_scoped_by_colon_boundary(self): DatabaseManager.reset_instance() db = DatabaseManager(db_url="sqlite:///:memory:") db.save_conversation_message("telegram_12345:chat", "user", "first user") db.save_conversation_message("telegram_123456:chat", "user", "second user") sessions = db.get_chat_sessions(session_prefix="telegram_12345") self.assertEqual(len(sessions), 1) self.assertEqual(sessions[0]["session_id"], "telegram_12345:chat") DatabaseManager.reset_instance() def test_get_chat_sessions_can_include_legacy_exact_session_id(self): DatabaseManager.reset_instance() db = DatabaseManager(db_url="sqlite:///:memory:") db.save_conversation_message("feishu_u1", "user", "legacy chat") db.save_conversation_message("feishu_u1:ask_600519", "user", "ask session") sessions = db.get_chat_sessions( session_prefix="feishu_u1:", extra_session_ids=["feishu_u1"], ) self.assertEqual({item["session_id"] for item in sessions}, {"feishu_u1", "feishu_u1:ask_600519"}) DatabaseManager.reset_instance() def test_conversation_summary_upsert_and_delete_with_session(self): DatabaseManager.reset_instance() db = DatabaseManager(db_url="sqlite:///:memory:") db.save_conversation_message("summary-session", "user", "hello") db.upsert_conversation_summary( "summary-session", "first summary", covered_message_id=1, source_message_count=1, estimated_tokens=10, ) db.upsert_conversation_summary( "summary-session", "updated summary", covered_message_id=2, source_message_count=2, estimated_tokens=12, ) summary = db.get_conversation_summary("summary-session") self.assertIsNotNone(summary) self.assertEqual(summary["summary"], "updated summary") self.assertEqual(summary["covered_message_id"], 2) self.assertEqual(summary["source_message_count"], 2) deleted = db.delete_conversation_session("summary-session") self.assertEqual(deleted, 1) self.assertIsNone(db.get_conversation_summary("summary-session")) DatabaseManager.reset_instance() def test_conversation_message_save_returns_id(self): DatabaseManager.reset_instance() db = DatabaseManager(db_url="sqlite:///:memory:") message_id = db.save_conversation_message("message-id-session", "user", "hello") self.assertIsInstance(message_id, int) self.assertGreater(message_id, 0) DatabaseManager.reset_instance() def test_conversation_user_turn_persists_and_updates_session_skills(self): DatabaseManager.reset_instance() db = DatabaseManager(db_url="sqlite:///:memory:") first_id = db.save_conversation_user_turn( "skill-session", "first question", ["technical", "risk"], ) second_id = db.save_conversation_user_turn( "skill-session", "use general analysis", [], ) self.assertGreater(first_id, 0) self.assertGreater(second_id, first_id) self.assertEqual( [message["content"] for message in db.get_conversation_messages("skill-session")], ["first question", "use general analysis"], ) self.assertEqual( db.get_conversation_session_selected_skill_ids("skill-session"), [], ) deleted = db.delete_conversation_session("skill-session") self.assertEqual(deleted, 2) self.assertIsNone( db.get_conversation_session_selected_skill_ids("skill-session") ) DatabaseManager.reset_instance() def test_conversation_user_turn_without_skill_update_keeps_session_state(self): DatabaseManager.reset_instance() db = DatabaseManager(db_url="sqlite:///:memory:") db.save_conversation_user_turn("skill-session", "first", ["technical"]) db.save_conversation_user_turn("skill-session", "follow up") self.assertEqual( db.get_conversation_session_selected_skill_ids("skill-session"), ["technical"], ) DatabaseManager.reset_instance() def test_conversation_user_turn_rolls_back_message_when_state_write_fails(self): DatabaseManager.reset_instance() db = DatabaseManager(db_url="sqlite:///:memory:") with patch("src.storage.sqlite_insert", side_effect=RuntimeError("state write failed")): with self.assertRaisesRegex(RuntimeError, "state write failed"): db.save_conversation_user_turn( "skill-session", "not accepted", ["technical"], ) self.assertEqual(db.get_conversation_messages("skill-session"), []) self.assertIsNone( db.get_conversation_session_selected_skill_ids("skill-session") ) DatabaseManager.reset_instance() def test_provider_turn_round_trip_preserves_protocol_fields_and_flags(self): DatabaseManager.reset_instance() db = DatabaseManager(db_url="sqlite:///:memory:") user_id = db.save_conversation_message("trace-session", "user", "question") assistant_id = db.save_conversation_message("trace-session", "assistant", "final") trace_messages = [ { "role": "assistant", "content": "checking", "reasoning_content": "reasoning", "tool_calls": [ { "id": "call_1", "name": "echo", "arguments": {"message": "hello"}, "provider_specific_fields": {"thought_signature": "sig"}, } ], }, {"role": "tool", "tool_call_id": "call_1", "content": "{\"ok\": true}"}, ] turn_id = db.save_agent_provider_turn( session_id="trace-session", run_id="run-1", provider="deepseek", model="deepseek/deepseek-chat", anchor_user_message_id=user_id, anchor_assistant_message_id=assistant_id, messages=trace_messages, contains_reasoning=True, contains_tool_calls=True, contains_thinking_blocks=False, must_roundtrip=True, estimated_tokens=42, ) rows = db.get_agent_provider_turns("trace-session") self.assertIsInstance(turn_id, int) self.assertEqual(len(rows), 1) self.assertEqual(rows[0]["messages"], trace_messages) self.assertTrue(rows[0]["contains_reasoning"]) self.assertTrue(rows[0]["contains_tool_calls"]) self.assertTrue(rows[0]["must_roundtrip"]) self.assertEqual(rows[0]["estimated_tokens"], 42) DatabaseManager.reset_instance() def test_provider_turns_do_not_appear_in_visible_or_web_messages_and_delete_with_session(self): DatabaseManager.reset_instance() db = DatabaseManager(db_url="sqlite:///:memory:") user_id = db.save_conversation_message("trace-hidden", "user", "visible question") assistant_id = db.save_conversation_message("trace-hidden", "assistant", "visible answer") db.save_agent_provider_turn( session_id="trace-hidden", run_id="run-hidden", provider="deepseek", model="deepseek/deepseek-chat", anchor_user_message_id=user_id, anchor_assistant_message_id=assistant_id, messages=[{"role": "assistant", "reasoning_content": "SECRET_REASONING", "tool_calls": []}], contains_reasoning=True, contains_tool_calls=True, contains_thinking_blocks=False, must_roundtrip=True, estimated_tokens=5, ) self.assertEqual( [(m["role"], m["content"]) for m in db.get_visible_conversation_messages("trace-hidden")], [("user", "visible question"), ("assistant", "visible answer")], ) self.assertEqual( [(m["role"], m["content"]) for m in db.get_conversation_history("trace-hidden")], [("user", "visible question"), ("assistant", "visible answer")], ) self.assertEqual( [(m["role"], m["content"]) for m in db.get_conversation_messages("trace-hidden")], [("user", "visible question"), ("assistant", "visible answer")], ) deleted = db.delete_conversation_session("trace-hidden") self.assertEqual(deleted, 2) self.assertEqual(db.get_agent_provider_turns("trace-hidden"), []) DatabaseManager.reset_instance() def test_provider_turn_retention_is_bucketed_by_session_provider_model(self): DatabaseManager.reset_instance() db = DatabaseManager(db_url="sqlite:///:memory:") for idx in range(5): user_id = db.save_conversation_message("retention", "user", f"q{idx}") assistant_id = db.save_conversation_message("retention", "assistant", f"a{idx}") db.save_agent_provider_turn( session_id="retention", run_id=f"run-{idx}", provider="deepseek", model="deepseek/deepseek-chat", anchor_user_message_id=user_id, anchor_assistant_message_id=assistant_id, messages=[{"role": "assistant", "reasoning_content": f"r{idx}", "tool_calls": [{"id": f"c{idx}", "name": "echo", "arguments": {}}]}], contains_reasoning=True, contains_tool_calls=True, contains_thinking_blocks=False, must_roundtrip=True, estimated_tokens=idx + 1, ) user_id = db.save_conversation_message("retention", "user", "other") assistant_id = db.save_conversation_message("retention", "assistant", "other") db.save_agent_provider_turn( session_id="retention", run_id="run-other", provider="anthropic", model="anthropic/claude-test", anchor_user_message_id=user_id, anchor_assistant_message_id=assistant_id, messages=[{"role": "assistant", "provider_blocks": [{"type": "thinking"}], "tool_calls": [{"id": "c-other", "name": "echo", "arguments": {}}]}], contains_reasoning=False, contains_tool_calls=True, contains_thinking_blocks=True, must_roundtrip=True, estimated_tokens=1, ) deepseek_rows = db.get_agent_provider_turns( "retention", provider="deepseek", model="deepseek/deepseek-chat", ) anthropic_rows = db.get_agent_provider_turns( "retention", provider="anthropic", model="anthropic/claude-test", ) self.assertEqual(len(deepseek_rows), 3) self.assertEqual([row["run_id"] for row in deepseek_rows], ["run-2", "run-3", "run-4"]) self.assertEqual(len(anthropic_rows), 1) DatabaseManager.reset_instance() def test_get_visible_conversation_messages_returns_ordered_visible_content(self): DatabaseManager.reset_instance() db = DatabaseManager(db_url="sqlite:///:memory:") db.save_conversation_message("visible-session", "system", "hidden") db.save_conversation_message("visible-session", "user", "question") db.save_conversation_message("visible-session", "assistant", "answer") messages = db.get_visible_conversation_messages("visible-session") self.assertEqual( [(item["role"], item["content"]) for item in messages], [("user", "question"), ("assistant", "answer")], ) self.assertIsInstance(messages[0]["id"], int) DatabaseManager.reset_instance() def test_get_visible_conversation_messages_limit_returns_ordered_tail(self): DatabaseManager.reset_instance() db = DatabaseManager(db_url="sqlite:///:memory:") for idx in range(25): db.save_conversation_message("visible-limit", "user", f"msg-{idx}") messages = db.get_visible_conversation_messages("visible-limit", limit=20) self.assertEqual(len(messages), 20) self.assertEqual(messages[0]["content"], "msg-5") self.assertEqual(messages[-1]["content"], "msg-24") DatabaseManager.reset_instance() def test_file_sqlite_enables_wal_and_busy_timeout(self): temp_dir = tempfile.TemporaryDirectory() db_path = os.path.join(temp_dir.name, "sqlite_pragmas.db") original_env = { "DATABASE_PATH": os.environ.get("DATABASE_PATH"), "SQLITE_BUSY_TIMEOUT_MS": os.environ.get("SQLITE_BUSY_TIMEOUT_MS"), "SQLITE_WAL_ENABLED": os.environ.get("SQLITE_WAL_ENABLED"), } try: os.environ["DATABASE_PATH"] = db_path os.environ["SQLITE_BUSY_TIMEOUT_MS"] = "1234" os.environ["SQLITE_WAL_ENABLED"] = "true" Config.reset_instance() DatabaseManager.reset_instance() db = DatabaseManager.get_instance() with db.get_session() as session: journal_mode = session.connection().exec_driver_sql("PRAGMA journal_mode").scalar() busy_timeout = session.connection().exec_driver_sql("PRAGMA busy_timeout").scalar() self.assertEqual(str(journal_mode).lower(), "wal") self.assertEqual(int(busy_timeout), 1234) finally: DatabaseManager.reset_instance() Config.reset_instance() for key, value in original_env.items(): if value is None: os.environ.pop(key, None) else: os.environ[key] = value temp_dir.cleanup() def test_get_instance_waits_for_cold_start_initialization(self): DatabaseManager.reset_instance() Config.reset_instance() temp_dir = tempfile.TemporaryDirectory() db_path = os.path.join(temp_dir.name, "sqlite_cold_start.db") original_database_path = os.environ.get("DATABASE_PATH") create_all_entered = threading.Event() competitor_entered = threading.Event() release_create_all = threading.Event() competitor_done = threading.Event() state_lock = threading.Lock() init_errors = [] competitor_errors = [] instances = [] query_values = [] original_create_all = Base.metadata.create_all def delayed_create_all(bind, *args, **kwargs): create_all_entered.set() if not release_create_all.wait(timeout=5): raise TimeoutError("Timed out waiting to release create_all") return original_create_all(bind, *args, **kwargs) def initialize_manager() -> None: try: db = DatabaseManager.get_instance() with state_lock: instances.append(db) except Exception as exc: with state_lock: init_errors.append(exc) def use_manager() -> None: try: competitor_entered.set() db = DatabaseManager.get_instance() session = db.get_session() try: value = session.connection().exec_driver_sql("SELECT 1").scalar() finally: session.close() with state_lock: instances.append(db) query_values.append(value) except Exception as exc: with state_lock: competitor_errors.append(exc) finally: competitor_done.set() try: os.environ["DATABASE_PATH"] = db_path Config.reset_instance() with patch.object(Base.metadata, "create_all", side_effect=delayed_create_all): init_thread = threading.Thread(target=initialize_manager) competitor_thread = threading.Thread(target=use_manager) init_thread.start() self.assertTrue(create_all_entered.wait(timeout=5)) competitor_thread.start() self.assertTrue(competitor_entered.wait(timeout=5)) self.assertFalse( competitor_done.wait(timeout=0.2), "DatabaseManager.get_instance() returned before initialization completed", ) release_create_all.set() init_thread.join(timeout=5) competitor_thread.join(timeout=5) self.assertFalse(init_thread.is_alive()) self.assertFalse(competitor_thread.is_alive()) self.assertEqual(init_errors, []) self.assertEqual(competitor_errors, []) self.assertEqual(query_values, [1]) self.assertEqual(len({id(instance) for instance in instances}), 1) finally: release_create_all.set() DatabaseManager.reset_instance() Config.reset_instance() if original_database_path is None: os.environ.pop("DATABASE_PATH", None) else: os.environ["DATABASE_PATH"] = original_database_path temp_dir.cleanup() def test_direct_construction_serializes_before_get_instance(self): DatabaseManager.reset_instance() Config.reset_instance() temp_dir = tempfile.TemporaryDirectory() direct_db_path = os.path.join(temp_dir.name, "direct.db") env_db_path = os.path.join(temp_dir.name, "env.db") direct_db_url = f"sqlite:///{direct_db_path}" original_database_path = os.environ.get("DATABASE_PATH") direct_init_entered = threading.Event() competitor_entered = threading.Event() allow_direct_init = threading.Event() competitor_done = threading.Event() state_lock = threading.Lock() errors = [] instances = [] query_values = [] original_init = DatabaseManager.__init__ def delayed_direct_init(self, db_url=None): if db_url == direct_db_url: direct_init_entered.set() if not competitor_entered.wait(timeout=5): raise TimeoutError("Timed out waiting for competitor") if not allow_direct_init.wait(timeout=5): raise TimeoutError("Timed out waiting to initialize direct instance") return original_init(self, db_url=db_url) def construct_directly() -> None: try: db = DatabaseManager(db_url=direct_db_url) with state_lock: instances.append(db) except Exception as exc: with state_lock: errors.append(exc) def use_get_instance() -> None: try: competitor_entered.set() db = DatabaseManager.get_instance() session = db.get_session() try: value = session.connection().exec_driver_sql("SELECT 1").scalar() finally: session.close() with state_lock: instances.append(db) query_values.append(value) except Exception as exc: with state_lock: errors.append(exc) finally: competitor_done.set() try: os.environ["DATABASE_PATH"] = env_db_path Config.reset_instance() with patch.object(DatabaseManager, "__init__", new=delayed_direct_init): direct_thread = threading.Thread(target=construct_directly) competitor_thread = threading.Thread(target=use_get_instance) direct_thread.start() self.assertTrue(direct_init_entered.wait(timeout=5)) competitor_thread.start() self.assertTrue(competitor_entered.wait(timeout=5)) self.assertFalse( competitor_done.wait(timeout=0.2), "get_instance() should not initialize over an in-flight direct construction", ) allow_direct_init.set() direct_thread.join(timeout=5) competitor_thread.join(timeout=5) self.assertFalse(direct_thread.is_alive()) self.assertFalse(competitor_thread.is_alive()) self.assertEqual(errors, []) self.assertEqual(query_values, [1]) self.assertEqual(len({id(instance) for instance in instances}), 1) self.assertEqual(DatabaseManager._instance._db_url, direct_db_url) finally: allow_direct_init.set() DatabaseManager.reset_instance() Config.reset_instance() if original_database_path is None: os.environ.pop("DATABASE_PATH", None) else: os.environ["DATABASE_PATH"] = original_database_path temp_dir.cleanup() def test_init_cleanup_preserves_original_initialization_error(self): DatabaseManager.reset_instance() original_error = RuntimeError("create all failed") cleanup_error = RuntimeError("dispose failed") def create_engine_with_failing_dispose(*args, **kwargs): engine = sqlalchemy_create_engine(*args, **kwargs) def failing_dispose() -> None: raise cleanup_error engine.dispose = failing_dispose return engine try: with patch("src.storage.create_engine", side_effect=create_engine_with_failing_dispose): with patch.object(Base.metadata, "create_all", side_effect=original_error): with self.assertRaisesRegex(RuntimeError, "create all failed") as ctx: DatabaseManager.get_instance() self.assertIs(ctx.exception, original_error) self.assertIsNone(DatabaseManager._instance) finally: DatabaseManager.reset_instance() def test_sqlite_write_transactions_begin_immediate(self): DatabaseManager.reset_instance() db = DatabaseManager(db_url="sqlite:///:memory:") session = db.get_session() connection = session.connection() try: with patch.object(db, "get_session", return_value=session): with patch.object(connection, "exec_driver_sql", wraps=connection.exec_driver_sql) as mock_exec: result = db._run_write_transaction("unit-test", lambda current_session: 7) self.assertEqual(result, 7) self.assertTrue( any(call.args == ("BEGIN IMMEDIATE",) for call in mock_exec.call_args_list) ) finally: DatabaseManager.reset_instance() def test_save_daily_data_sqlite_concurrent_same_code_date_counts_only_new_rows(self): DatabaseManager.reset_instance() temp_dir = tempfile.TemporaryDirectory() db_path = os.path.join(temp_dir.name, "sqlite_daily_concurrency.db") db = DatabaseManager(db_url=f"sqlite:///{db_path}") results = [] results_lock = threading.Lock() start_barrier = threading.Barrier(2) def worker() -> None: start_barrier.wait() count = db.save_daily_data( pd.DataFrame( [ { 'date': date(2026, 4, 1), 'open': 10, 'high': 11, 'low': 9, 'close': 10.5, 'volume': 100, 'amount': 1050, 'pct_chg': 1.2, 'ma5': 10.1, 'ma10': 10.2, 'ma20': 10.3, 'volume_ratio': 1.0, } ] ), code='600519', data_source='test', ) with results_lock: results.append(count) threads = [threading.Thread(target=worker) for _ in range(2)] for thread in threads: thread.start() for thread in threads: thread.join() try: self.assertCountEqual(results, [1, 0]) with db.get_session() as session: total = session.execute( select(func.count()).select_from(StockDaily).where( and_( StockDaily.code == '600519', StockDaily.date == date(2026, 4, 1), ) ) ).scalar() self.assertEqual(total, 1) finally: temp_dir.cleanup() DatabaseManager.reset_instance() # ------------------------------------------------------------------ # Story 1.2 — canonical_id column + dual-write (issue #2207 PR2) # ------------------------------------------------------------------ @staticmethod def _create_legacy_stock_daily_without_canonical_id(db_path: str) -> None: """Pre-PR2 ``stock_daily`` schema: no ``canonical_id`` column.""" with sqlite3.connect(db_path) as conn: conn.execute( """CREATE TABLE stock_daily ( id INTEGER PRIMARY KEY AUTOINCREMENT, code VARCHAR(10) NOT NULL, date DATE NOT NULL, open FLOAT, high FLOAT, low FLOAT, close FLOAT, volume FLOAT, amount FLOAT, pct_chg FLOAT, ma5 FLOAT, ma10 FLOAT, ma20 FLOAT, volume_ratio FLOAT, data_source VARCHAR(50), created_at DATETIME, updated_at DATETIME )""" ) conn.execute( "CREATE UNIQUE INDEX uix_code_date ON stock_daily (code, date)" ) conn.execute( "CREATE INDEX ix_code_date ON stock_daily (code, date)" ) @staticmethod def _create_legacy_stock_daily_with_canonical_id(db_path: str) -> None: """``stock_daily`` schema WITH a ``canonical_id`` column (for repair tests).""" with sqlite3.connect(db_path) as conn: conn.execute( """CREATE TABLE stock_daily ( id INTEGER PRIMARY KEY AUTOINCREMENT, code VARCHAR(10) NOT NULL, date DATE NOT NULL, open FLOAT, high FLOAT, low FLOAT, close FLOAT, volume FLOAT, amount FLOAT, pct_chg FLOAT, ma5 FLOAT, ma10 FLOAT, ma20 FLOAT, volume_ratio FLOAT, data_source VARCHAR(50), created_at DATETIME, updated_at DATETIME, canonical_id VARCHAR(32) )""" ) conn.execute( "CREATE UNIQUE INDEX uix_code_date ON stock_daily (code, date)" ) conn.execute( "CREATE INDEX ix_code_date ON stock_daily (code, date)" ) @staticmethod def _make_temp_db_path() -> tuple: """Return (db_dir, db_path); Windows-safe cleanup via shutil.rmtree.""" db_dir = tempfile.mkdtemp(prefix="dsa_canonical_id_") db_path = os.path.join(db_dir, "test.db") return db_dir, db_path @staticmethod def _cleanup_temp_dir(db_dir: str) -> None: # Windows: the SQLAlchemy engine pool may still hold the SQLite file # handle briefly after dispose(); ignore_errors lets teardown succeed. shutil.rmtree(db_dir, ignore_errors=True) def test_canonical_id_migration_adds_column_and_backfills_existing_rows(self): """AC 1: ALTER adds ``canonical_id``; existing rows backfilled via parser.""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() try: self._create_legacy_stock_daily_without_canonical_id(db_path) with sqlite3.connect(db_path) as conn: conn.executemany( """INSERT INTO stock_daily ( code, date, close, data_source ) VALUES (?, ?, ?, ?)""", [ ("sh000300", "2026-01-01", 4000.0, "legacy"), ("600519", "2026-01-01", 1600.0, "legacy"), ("AAPL", "2026-01-01", 150.0, "legacy"), ], ) DatabaseManager(db_url=f"sqlite:///{db_path}") with sqlite3.connect(db_path) as conn: columns = {row[1] for row in conn.execute("PRAGMA table_info(stock_daily)").fetchall()} rows = { code: canonical_id for code, canonical_id in conn.execute( "SELECT code, canonical_id FROM stock_daily ORDER BY id" ).fetchall() } self.assertIn("canonical_id", columns) # sh000300 is a registered index alias → sh000300 self.assertEqual(rows["sh000300"], "sh000300") # 600519 bare A-share → sh600519 (contract: bare codes default stock, # classifier routes 6-prefixed to SH) self.assertEqual(rows["600519"], "sh600519") # AAPL bare US ticker → AAPL (canonical US form, no prefix) self.assertEqual(rows["AAPL"], "AAPL") finally: DatabaseManager.reset_instance() Config.reset_instance() self._cleanup_temp_dir(db_dir) def test_canonical_id_migration_is_idempotent(self): """AC 2: second startup detects the column and skips ALTER; backfill is a no-op.""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() try: self._create_legacy_stock_daily_without_canonical_id(db_path) with sqlite3.connect(db_path) as conn: conn.execute( "INSERT INTO stock_daily (code, date, close) VALUES (?, ?, ?)", ("600519", "2026-01-01", 1600.0), ) # First run: adds column + backfills. DatabaseManager(db_url=f"sqlite:///{db_path}") with sqlite3.connect(db_path) as conn: first_pass = conn.execute( "SELECT canonical_id FROM stock_daily WHERE code='600519'" ).fetchone()[0] self.assertEqual(first_pass, "sh600519") DatabaseManager.reset_instance() # Second run: must not raise, must not duplicate the column. DatabaseManager(db_url=f"sqlite:///{db_path}") with sqlite3.connect(db_path) as conn: columns = [row[1] for row in conn.execute("PRAGMA table_info(stock_daily)").fetchall()] second_pass = conn.execute( "SELECT canonical_id FROM stock_daily WHERE code='600519'" ).fetchone()[0] self.assertEqual(columns.count("canonical_id"), 1) self.assertEqual(second_pass, "sh600519") finally: DatabaseManager.reset_instance() Config.reset_instance() self._cleanup_temp_dir(db_dir) def test_canonical_id_migration_on_fresh_empty_db(self): """AC: empty/new DB has the column from create_all; backfill finds 0 rows.""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() try: DatabaseManager(db_url=f"sqlite:///{db_path}") with sqlite3.connect(db_path) as conn: columns = {row[1] for row in conn.execute("PRAGMA table_info(stock_daily)").fetchall()} count = conn.execute("SELECT COUNT(*) FROM stock_daily").fetchone()[0] self.assertIn("canonical_id", columns) self.assertEqual(count, 0) indexes = self._list_sqlite_indexes(db_path, "stock_daily") self.assertEqual(indexes.get("ix_stock_daily_canonical_id"), ["canonical_id"]) finally: DatabaseManager.reset_instance() Config.reset_instance() self._cleanup_temp_dir(db_dir) def test_canonical_id_migration_creates_plain_non_unique_index(self): """AC 6/9: index is plain (not unique) so alias rows can coexist.""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() try: DatabaseManager(db_url=f"sqlite:///{db_path}") unique_indexes = self._list_sqlite_unique_indexes(db_path, "stock_daily") self.assertNotIn("ix_stock_daily_canonical_id", unique_indexes) indexes = self._list_sqlite_indexes(db_path, "stock_daily") self.assertEqual(indexes.get("ix_stock_daily_canonical_id"), ["canonical_id"]) finally: DatabaseManager.reset_instance() Config.reset_instance() self._cleanup_temp_dir(db_dir) def test_canonical_id_backfill_bare_code_colliding_with_index_stays_stock(self): """``_derive_canonical_id`` no longer reads ``matched_index.canonical_id``. A bare ``000300`` (which collides with the CSI-300 index) now derives to the stock-path canonical_id ``sz000300`` — the parser contract says bare codes are always stock. The index conflict is surfaced via ``matched_index`` only, never used to override the canonical_id.""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() try: self._create_legacy_stock_daily_without_canonical_id(db_path) with sqlite3.connect(db_path) as conn: conn.execute( "INSERT INTO stock_daily (code, date, close) VALUES (?, ?, ?)", ("000300", "2026-01-01", 4000.0), ) DatabaseManager(db_url=f"sqlite:///{db_path}") with sqlite3.connect(db_path) as conn: canonical_id = conn.execute( "SELECT canonical_id FROM stock_daily WHERE code='000300'" ).fetchone()[0] # Bare code resolves to the stock-path canonical_id. self.assertEqual(canonical_id, "sz000300") finally: DatabaseManager.reset_instance() Config.reset_instance() self._cleanup_temp_dir(db_dir) def test_canonical_id_backfill_bare_non_index_code_stays_stock(self): """Index-aware backfill does not affect bare non-index codes: ``600519`` has no registry hit (``matched_index is None``) and backfills to the stock-path canonical_id ``sh600519`` (contract #2 — bare → stock).""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() try: self._create_legacy_stock_daily_without_canonical_id(db_path) with sqlite3.connect(db_path) as conn: conn.execute( "INSERT INTO stock_daily (code, date, close) VALUES (?, ?, ?)", ("600519", "2026-01-01", 1600.0), ) DatabaseManager(db_url=f"sqlite:///{db_path}") with sqlite3.connect(db_path) as conn: canonical_id = conn.execute( "SELECT canonical_id FROM stock_daily WHERE code='600519'" ).fetchone()[0] self.assertEqual(canonical_id, "sh600519") finally: DatabaseManager.reset_instance() Config.reset_instance() self._cleanup_temp_dir(db_dir) def test_canonical_id_backfill_preserves_duplicate_alias_rows(self): """AC 9: multiple rows sharing canonical_id + date coexist (plain index).""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() try: self._create_legacy_stock_daily_without_canonical_id(db_path) with sqlite3.connect(db_path) as conn: # Two distinct code aliases for the same underlying, same date. conn.executemany( """INSERT INTO stock_daily (code, date, close) VALUES (?, ?, ?)""", [ ("sh000300", "2026-01-01", 4000.0), ("000300.SH", "2026-01-01", 4001.0), ], ) DatabaseManager(db_url=f"sqlite:///{db_path}") with sqlite3.connect(db_path) as conn: rows = conn.execute( "SELECT code, canonical_id FROM stock_daily ORDER BY id" ).fetchall() self.assertEqual(len(rows), 2) # Both survive — the plain index allows duplicate canonical_ids. canonical_ids = {row[1] for row in rows} self.assertEqual(canonical_ids, {"sh000300"}) finally: DatabaseManager.reset_instance() Config.reset_instance() self._cleanup_temp_dir(db_dir) def test_canonical_id_backfill_survives_per_row_derivation_failure(self): """D1 degrade: one unparseable row must not trap the batch scan loop. The backfill advances an id cursor, so a row whose derivation raises is skipped past (left NULL, retried next startup) instead of being re-selected forever. Good rows after the bad one must still backfill. """ real_parse = stock_list_parser_module.parse_analysis_target def flaky_parse(code): if code == "!!broken!!": raise RuntimeError("parser boom") return real_parse(code) DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() try: self._create_legacy_stock_daily_without_canonical_id(db_path) with sqlite3.connect(db_path) as conn: # Bad row FIRST, good rows after — maximizes re-select exposure # if the scan ever regresses to "re-select NULL rows". conn.executemany( """INSERT INTO stock_daily (code, date, close) VALUES (?, ?, ?)""", [ ("!!broken!!", "2026-01-01", 1.0), ("sh600519", "2026-01-01", 1700.0), ("AAPL", "2026-01-01", 210.0), ], ) with patch( "src.services.stock_list_parser.parse_analysis_target", side_effect=flaky_parse, ): DatabaseManager(db_url=f"sqlite:///{db_path}") with sqlite3.connect(db_path) as conn: rows = dict( conn.execute( "SELECT code, canonical_id FROM stock_daily" ).fetchall() ) self.assertIsNone(rows["!!broken!!"]) self.assertEqual(rows["sh600519"], "sh600519") self.assertEqual(rows["AAPL"], "AAPL") finally: DatabaseManager.reset_instance() Config.reset_instance() self._cleanup_temp_dir(db_dir) def test_canonical_id_backfill_skips_empty_derived_values(self): """Review fix: parser returning None/'' (without raising) must not count as backfilled nor persist an empty string as a stable key.""" real_parse = stock_list_parser_module.parse_analysis_target def empty_parse(code): if code == "!!empty!!": return None return real_parse(code) DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() try: self._create_legacy_stock_daily_without_canonical_id(db_path) with sqlite3.connect(db_path) as conn: conn.executemany( """INSERT INTO stock_daily (code, date, close) VALUES (?, ?, ?)""", [ ("!!empty!!", "2026-01-01", 1.0), ("sh600519", "2026-01-01", 1700.0), ], ) with patch( "src.services.stock_list_parser.parse_analysis_target", side_effect=empty_parse, ): DatabaseManager(db_url=f"sqlite:///{db_path}") with sqlite3.connect(db_path) as conn: rows = dict( conn.execute( "SELECT code, canonical_id FROM stock_daily" ).fetchall() ) self.assertIsNone(rows["!!empty!!"]) self.assertEqual(rows["sh600519"], "sh600519") finally: DatabaseManager.reset_instance() Config.reset_instance() self._cleanup_temp_dir(db_dir) def test_save_daily_data_dual_writes_explicit_canonical_id(self): """AC 3: explicit ``canonical_id`` is written to both columns on upsert.""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() db = DatabaseManager(db_url=f"sqlite:///{db_path}") try: df = pd.DataFrame( [ { 'date': date(2026, 4, 1), 'open': 10, 'high': 11, 'low': 9, 'close': 10.5, 'volume': 100, 'amount': 1050, 'pct_chg': 1.2, 'ma5': 10.1, 'ma10': 10.2, 'ma20': 10.3, 'volume_ratio': 1.0, } ] ) db.save_daily_data( df, code="sh000016", data_source="tencent", canonical_id="sh000016", ) with db.get_session() as session: row = session.execute( select(StockDaily).where( and_(StockDaily.code == "sh000016", StockDaily.date == date(2026, 4, 1)) ) ).scalar_one() self.assertEqual(row.code, "sh000016") self.assertEqual(row.canonical_id, "sh000016") # Upsert path: re-saving updates canonical_id too. db.save_daily_data( df, code="sh000016", data_source="tencent", canonical_id="sh000016-v2", ) with db.get_session() as session: row = session.execute( select(StockDaily).where( and_(StockDaily.code == "sh000016", StockDaily.date == date(2026, 4, 1)) ) ).scalar_one() self.assertEqual(row.canonical_id, "sh000016-v2") finally: DatabaseManager.reset_instance() self._cleanup_temp_dir(db_dir) def test_save_daily_data_derives_canonical_id_when_omitted(self): """D1/AC: ``canonical_id`` omitted → auto-derive via parser; parser failure → NULL, no raise.""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() db = DatabaseManager(db_url=f"sqlite:///{db_path}") try: df = pd.DataFrame( [ { 'date': date(2026, 4, 1), 'open': 10, 'high': 11, 'low': 9, 'close': 10.5, 'volume': 100, 'amount': 1050, 'pct_chg': 1.2, 'ma5': 10.1, 'ma10': 10.2, 'ma20': 10.3, 'volume_ratio': 1.0, } ] ) # Main path: bare A-share → sh600519. db.save_daily_data(df, code="600519", data_source="test") with db.get_session() as session: row = session.execute( select(StockDaily).where( and_(StockDaily.code == "600519", StockDaily.date == date(2026, 4, 1)) ) ).scalar_one() self.assertEqual(row.canonical_id, "sh600519") # Degraded path: parser failure on re-save must NOT overwrite a # previously backfilled non-NULL canonical_id with NULL (review # fix: coalesce preserves the existing value). No exception # bubbles out of save_daily_data (D1). with patch( "src.services.stock_list_parser.parse_analysis_target", side_effect=RuntimeError("parser boom"), ): db.save_daily_data(df, code="600519", data_source="test") with db.get_session() as session: row = session.execute( select(StockDaily).where( and_(StockDaily.code == "600519", StockDaily.date == date(2026, 4, 1)) ) ).scalar_one() self.assertEqual(row.canonical_id, "sh600519") finally: DatabaseManager.reset_instance() self._cleanup_temp_dir(db_dir) def test_save_daily_data_empty_string_canonical_id_triggers_derivation(self): """Review fix: empty-string canonical_id is treated like None — derived via parser, never persisted as an empty key.""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() db = DatabaseManager(db_url=f"sqlite:///{db_path}") try: df = pd.DataFrame( [ { 'date': date(2026, 4, 3), 'open': 10, 'high': 11, 'low': 9, 'close': 10.5, 'volume': 100, 'amount': 1050, 'pct_chg': 1.2, 'ma5': 10.1, 'ma10': 10.2, 'ma20': 10.3, 'volume_ratio': 1.0, } ] ) db.save_daily_data(df, code="600519", data_source="test", canonical_id="") with db.get_session() as session: row = session.execute( select(StockDaily).where( and_(StockDaily.code == "600519", StockDaily.date == date(2026, 4, 3)) ) ).scalar_one() self.assertEqual(row.canonical_id, "sh600519") finally: DatabaseManager.reset_instance() self._cleanup_temp_dir(db_dir) def test_save_daily_data_first_insert_parser_failure_writes_null(self): """D1: first insert with parser failure → row created with NULL canonical_id, no raise.""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() db = DatabaseManager(db_url=f"sqlite:///{db_path}") try: df = pd.DataFrame( [ { 'date': date(2026, 4, 2), 'open': 10, 'high': 11, 'low': 9, 'close': 10.5, 'volume': 100, 'amount': 1050, 'pct_chg': 1.2, 'ma5': 10.1, 'ma10': 10.2, 'ma20': 10.3, 'volume_ratio': 1.0, } ] ) with patch( "src.services.stock_list_parser.parse_analysis_target", side_effect=RuntimeError("parser boom"), ): db.save_daily_data(df, code="600519", data_source="test") with db.get_session() as session: row = session.execute( select(StockDaily).where( and_(StockDaily.code == "600519", StockDaily.date == date(2026, 4, 2)) ) ).scalar_one() self.assertIsNone(row.canonical_id) finally: DatabaseManager.reset_instance() self._cleanup_temp_dir(db_dir) def test_save_daily_data_derives_stock_canonical_id_for_bare_index_collision(self): """``save_daily_data(df, code="000300")`` with no explicit canonical_id writes ``sz000300`` (stock-path canonical_id), NOT ``sh000300`` (the index canonical_id). Bare codes are always stock.""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() db = DatabaseManager(db_url=f"sqlite:///{db_path}") try: df = pd.DataFrame( [ { 'date': date(2026, 4, 5), 'open': 10, 'high': 11, 'low': 9, 'close': 10.5, 'volume': 100, 'amount': 1050, 'pct_chg': 1.2, 'ma5': 10.1, 'ma10': 10.2, 'ma20': 10.3, 'volume_ratio': 1.0, } ] ) db.save_daily_data(df, code="000300", data_source="test") with db.get_session() as session: row = session.execute( select(StockDaily).where( and_(StockDaily.code == "000300", StockDaily.date == date(2026, 4, 5)) ) ).scalar_one() self.assertEqual(row.canonical_id, "sz000300") finally: DatabaseManager.reset_instance() self._cleanup_temp_dir(db_dir) def test_save_daily_data_explicit_index_forms_derive_index_canonical_id(self): """Explicit index forms (``sh000300`` / ``000300.SH``) derive to the index canonical_id ``sh000300``, while the bare code ``000300`` derives to the stock-path ``sz000300``. Explicit index and bare stock are intentionally different buckets.""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() db = DatabaseManager(db_url=f"sqlite:///{db_path}") try: df = pd.DataFrame( [ { 'date': date(2026, 4, 6), 'open': 10, 'high': 11, 'low': 9, 'close': 10.5, 'volume': 100, 'amount': 1050, 'pct_chg': 1.2, 'ma5': 10.1, 'ma10': 10.2, 'ma20': 10.3, 'volume_ratio': 1.0, } ] ) for code in ("000300", "sh000300", "000300.SH"): db.save_daily_data(df, code=code, data_source="test") with db.get_session() as session: rows = session.execute( select(StockDaily).order_by(StockDaily.id) ).scalars().all() self.assertEqual(len(rows), 3) self.assertEqual( {row.code: row.canonical_id for row in rows}, { "000300": "sz000300", "sh000300": "sh000300", "000300.SH": "sh000300", }, ) finally: DatabaseManager.reset_instance() self._cleanup_temp_dir(db_dir) def test_canonical_id_migration_raises_when_column_inspection_fails(self): """AC: inspect() failure on stock_daily must raise (no silent downgrade). Mirrors the decision_profile inspection-failure test (L417-438). The patched inspector only breaks for ``stock_daily`` so the earlier ``decision_signal_profile`` migration (which runs first in the ``__init__`` ensure chain) still passes and execution reaches the ``_ensure_stock_daily_canonical_id`` step. """ real_inspector = inspect class BrokenInspector: def __init__(self, engine): self._real = real_inspector(engine) def has_table(self, table_name: str) -> bool: return self._real.has_table(table_name) def get_columns(self, table_name: str): if table_name == StockDaily.__tablename__: raise RuntimeError("inspection failed") return self._real.get_columns(table_name) def get_indexes(self, table_name: str): return self._real.get_indexes(table_name) DatabaseManager.reset_instance() try: with patch("src.storage.inspect", side_effect=BrokenInspector): with self.assertLogs("src.storage", level="ERROR") as logs: with self.assertRaises(RuntimeError): DatabaseManager(db_url="sqlite:///:memory:") self.assertIn( "canonical_id migration cannot continue safely", "\n".join(logs.output), ) finally: DatabaseManager.reset_instance() Config.reset_instance() # ------------------------------------------------------------------ # canonical_id repair # ------------------------------------------------------------------ def test_derive_canonical_id_bare_conflict_stays_stock(self): """``_derive_canonical_id`` returns the parser stock canonical for bare conflict codes, never the index canonical.""" DatabaseManager.reset_instance() db = DatabaseManager(db_url="sqlite:///:memory:") try: self.assertEqual(db._derive_canonical_id("000016"), "sz000016") self.assertEqual(db._derive_canonical_id("930955"), "bj930955") self.assertEqual(db._derive_canonical_id("sh000016"), "sh000016") self.assertEqual(db._derive_canonical_id("930955.CSI"), "csi930955") finally: DatabaseManager.reset_instance() def test_derive_canonical_id_rejects_unregistered_csi(self): """PR #2267 review fix: an unsupported identity (unregistered ``csi`` prefix or ``.CSI`` suffix) must NOT enter a persistent canonical bucket — ``_derive_canonical_id`` returns None so the caller persists NULL instead of the raw token.""" DatabaseManager.reset_instance() db = DatabaseManager(db_url="sqlite:///:memory:") try: self.assertIsNone(db._derive_canonical_id("csi930956")) self.assertIsNone(db._derive_canonical_id("CSI930956")) self.assertIsNone(db._derive_canonical_id("930956.CSI")) self.assertIsNone(db._derive_canonical_id("csi000300")) finally: DatabaseManager.reset_instance() def test_canonical_id_repair_fixes_bare_misbucketed_rows(self): """Rows whose bare code has an erroneous index canonical are repaired to the parser stock canonical; explicit index rows and correct stock rows are untouched.""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() try: self._create_legacy_stock_daily_with_canonical_id(db_path) with sqlite3.connect(db_path) as conn: conn.executemany( """INSERT INTO stock_daily (code, date, close, canonical_id) VALUES (?, ?, ?, ?)""", [ # Mis-bucketed bare codes (old index-aware derivation). ("000001", "2026-01-01", 1.0, "sh000001"), ("000016", "2026-01-02", 2.0, "sh000016"), ("000688", "2026-01-03", 3.0, "sh000688"), ("930955", "2026-01-04", 4.0, "csi930955"), # Explicit index rows — must NOT be modified. ("sh000016", "2026-01-05", 5.0, "sh000016"), ("930955.CSI", "2026-01-06", 6.0, "csi930955"), # Correct stock row — must NOT be modified. ("600519", "2026-01-07", 7.0, "sh600519"), # Unrelated row — must NOT be modified. ("AAPL", "2026-01-08", 8.0, "AAPL"), ], ) DatabaseManager(db_url=f"sqlite:///{db_path}") with sqlite3.connect(db_path) as conn: rows = dict( conn.execute( "SELECT code, canonical_id FROM stock_daily ORDER BY id" ).fetchall() ) self.assertEqual(rows["000001"], "sz000001") self.assertEqual(rows["000016"], "sz000016") self.assertEqual(rows["000688"], "sz000688") self.assertEqual(rows["930955"], "bj930955") # Explicit index rows preserved. self.assertEqual(rows["sh000016"], "sh000016") self.assertEqual(rows["930955.CSI"], "csi930955") # Correct stock + unrelated rows preserved. self.assertEqual(rows["600519"], "sh600519") self.assertEqual(rows["AAPL"], "AAPL") finally: DatabaseManager.reset_instance() Config.reset_instance() self._cleanup_temp_dir(db_dir) def test_canonical_id_repair_is_idempotent(self): """Running the repair a second time repairs 0 rows.""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() try: self._create_legacy_stock_daily_with_canonical_id(db_path) with sqlite3.connect(db_path) as conn: conn.execute( "INSERT INTO stock_daily (code, date, close, canonical_id) " "VALUES (?, ?, ?, ?)", ("000016", "2026-01-01", 2.0, "sh000016"), ) DatabaseManager(db_url=f"sqlite:///{db_path}") with sqlite3.connect(db_path) as conn: first = conn.execute( "SELECT canonical_id FROM stock_daily WHERE code='000016'" ).fetchone()[0] self.assertEqual(first, "sz000016") DatabaseManager.reset_instance() with self.assertLogs("src.storage", level="INFO") as logs: DatabaseManager(db_url=f"sqlite:///{db_path}") log_text = "\n".join(logs.output) self.assertIn("repaired_count=0", log_text) finally: DatabaseManager.reset_instance() Config.reset_instance() self._cleanup_temp_dir(db_dir) def test_canonical_id_repair_skips_when_registry_empty(self): """When the index registry is empty the repair is a no-op and logs a WARNING; no rows are modified.""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() try: self._create_legacy_stock_daily_with_canonical_id(db_path) with sqlite3.connect(db_path) as conn: conn.execute( "INSERT INTO stock_daily (code, date, close, canonical_id) " "VALUES (?, ?, ?, ?)", ("000016", "2026-01-01", 2.0, "sh000016"), ) with patch( "src.data.stock_index_loader._load_active_index_rows", return_value=[], ), self.assertLogs("src.storage", level="WARNING") as logs: DatabaseManager(db_url=f"sqlite:///{db_path}") with sqlite3.connect(db_path) as conn: canonical_id = conn.execute( "SELECT canonical_id FROM stock_daily WHERE code='000016'" ).fetchone()[0] self.assertEqual(canonical_id, "sh000016") self.assertTrue( any("registry is empty" in record.getMessage() for record in logs.records) ) finally: DatabaseManager.reset_instance() Config.reset_instance() self._cleanup_temp_dir(db_dir) def test_canonical_id_repair_handles_concurrent_rewrite_safely(self): """A conditional UPDATE that loses the race (row already rewritten) is skipped, not double-counted.""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() try: self._create_legacy_stock_daily_with_canonical_id(db_path) with sqlite3.connect(db_path) as conn: conn.execute( "INSERT INTO stock_daily (code, date, close, canonical_id) " "VALUES (?, ?, ?, ?)", ("000016", "2026-01-01", 2.0, "sh000016"), ) real_parse = stock_list_parser_module.parse_analysis_target def racing_parse(code): # Simulate a concurrent writer that already fixed the row. if code == "000016": with sqlite3.connect(db_path) as conn: conn.execute( "UPDATE stock_daily SET canonical_id='sz000016' " "WHERE code='000016'" ) conn.commit() return real_parse(code) with patch( "src.services.stock_list_parser.parse_analysis_target", side_effect=racing_parse, ): DatabaseManager(db_url=f"sqlite:///{db_path}") with sqlite3.connect(db_path) as conn: canonical_id = conn.execute( "SELECT canonical_id FROM stock_daily WHERE code='000016'" ).fetchone()[0] self.assertEqual(canonical_id, "sz000016") finally: DatabaseManager.reset_instance() Config.reset_instance() self._cleanup_temp_dir(db_dir) def test_canonical_id_repair_queries_only_relevant_rows(self): """Gap 6: the repair SQL filters to rows whose canonical_id is in the active index set, so rows with unrelated canonical_ids are never scanned/derived.""" DatabaseManager.reset_instance() db_dir, db_path = self._make_temp_db_path() try: self._create_legacy_stock_daily_with_canonical_id(db_path) with sqlite3.connect(db_path) as conn: conn.executemany( """INSERT INTO stock_daily (code, date, close, canonical_id) VALUES (?, ?, ?, ?)""", [ # Mis-bucketed bare code — must be repaired. ("000016", "2026-01-01", 2.0, "sh000016"), # Unrelated canonical — must NOT be touched. ("600519", "2026-01-02", 7.0, "sh600519"), ("AAPL", "2026-01-03", 8.0, "AAPL"), # Explicit index row — must NOT be touched. ("sh000016", "2026-01-04", 5.0, "sh000016"), ], ) DatabaseManager(db_url=f"sqlite:///{db_path}") with sqlite3.connect(db_path) as conn: rows = dict( conn.execute( "SELECT code, canonical_id FROM stock_daily ORDER BY id" ).fetchall() ) # Only the mis-bucketed bare code is repaired. self.assertEqual(rows["000016"], "sz000016") self.assertEqual(rows["600519"], "sh600519") self.assertEqual(rows["AAPL"], "AAPL") self.assertEqual(rows["sh000016"], "sh000016") finally: DatabaseManager.reset_instance() Config.reset_instance() self._cleanup_temp_dir(db_dir) if __name__ == '__main__': unittest.main()