from __future__ import annotations import tempfile import unittest import sqlite3 from pathlib import Path from database import ReviewDatabase from security import hash_password, verify_password class AccountAccessTests(unittest.TestCase): def setUp(self): self.temp = tempfile.TemporaryDirectory() self.database = ReviewDatabase(Path(self.temp.name) / "review.db") def tearDown(self): self.temp.cleanup() def test_first_user_is_admin_and_following_users_are_regular(self): first = self.database.create_user("admin_user", "salt", "hash") second = self.database.create_user("member_user", "salt", "hash") self.assertEqual(first["role"], "admin") self.assertEqual(second["role"], "user") self.assertEqual(self.database.user_access(first["id"])["role"], "admin") self.assertEqual(self.database.user_access(second["id"])["role"], "user") def test_admin_role_can_also_hold_an_explicit_membership(self): admin = self.database.create_user("admin_member", "salt", "hash") updated = self.database.update_membership( admin["id"], "active", "内部会员", "2026-07-22T00:00:00+00:00", "2026-08-23T00:00:00+00:00", ) access = self.database.user_access(admin["id"]) self.assertTrue(updated) self.assertEqual(access["role"], "admin") self.assertEqual(access["membership_status"], "active") def test_membership_mode_system_settings_and_usage_are_persistent(self): user = self.database.create_user("member_user", "salt", "hash") self.database.update_user_llm_mode(user["id"], "platform") updated = self.database.update_membership( user["id"], "active", "内部会员", "2026-07-22T00:00:00+00:00", "2026-08-23T00:00:00+00:00", ) self.database.save_system_setting("credentials", "encrypted") self.database.record_llm_usage( user["id"], "mentor", "platform", "model", "success", 1200 ) access = self.database.user_access(user["id"]) self.assertTrue(updated) self.assertEqual(access["llm_mode"], "platform") self.assertEqual(access["membership_status"], "active") self.assertEqual(access["membership_plan"], "内部会员") self.assertEqual(self.database.get_system_setting("credentials"), "encrypted") self.assertEqual( self.database.count_llm_usage_since( user["id"], "platform", "2026-01-01T00:00:00+00:00" ), 1, ) def test_password_can_be_rotated_without_changing_account_access(self): old_salt, old_hash = hash_password("OldPassword123") user = self.database.create_user("password_user", old_salt, old_hash) new_salt, new_hash = hash_password("NewPassword456") self.assertTrue( self.database.update_user_password(user["id"], new_salt, new_hash) ) stored = self.database.user_password(user["id"]) self.assertFalse( verify_password("OldPassword123", stored["password_salt"], stored["password_hash"]) ) self.assertTrue( verify_password("NewPassword456", stored["password_salt"], stored["password_hash"]) ) self.assertEqual(self.database.user_access(user["id"])["role"], "admin") def test_latest_real_snapshot_skips_demo_and_supports_strict_previous_date(self): self.database.save_snapshot( "20260720", "tushare", {"meta": {"trade_date": "2026-07-20", "source": "tushare"}} ) self.database.save_snapshot( "20260721", "demo", {"meta": {"trade_date": "2026-07-21", "source": "demo"}} ) latest = self.database.get_latest_real_snapshot("20260722") previous = self.database.get_latest_real_snapshot("20260721", strictly_before=True) self.assertEqual(latest["meta"]["trade_date"], "2026-07-20") self.assertEqual(previous["meta"]["trade_date"], "2026-07-20") def test_stock_master_search_supports_exact_name_and_code(self): self.database.upsert_stock_master( [ { "ts_code": "002141.SZ", "name": "贤丰控股", "industry": "元件", "market": "主板", "list_date": "20071228", } ] ) self.assertEqual(self.database.search_stock_master("贤丰控股")[0]["code"], "002141") self.assertEqual(self.database.search_stock_master("002141")[0]["name"], "贤丰控股") def test_review_notes_are_scoped_to_their_owner(self): first = self.database.create_user("note_owner", "salt", "hash") second = self.database.create_user("other_reader", "salt", "hash") note_id = self.database.save_note( first["id"], "002141", "贤丰控股", "20260721", "只属于甲", "明日观察" ) self.assertEqual(len(self.database.list_notes(first["id"], code="002141")), 1) self.assertEqual(self.database.list_notes(second["id"], code="002141"), []) with self.assertRaises(ValueError): self.database.save_note( second["id"], "002141", "贤丰控股", "20260721", "越权修改", "", note_id, ) self.assertFalse(self.database.delete_note(second["id"], note_id)) self.assertTrue(self.database.delete_note(first["id"], note_id)) def test_watchlist_is_scoped_to_its_owner(self): first = self.database.create_user("watch_owner", "salt", "hash") second = self.database.create_user("other_watcher", "salt", "hash") self.database.save_watchlist(first["id"], "002141", "贤丰控股", "元件", "red") self.database.save_watchlist(second["id"], "002141", "贤丰控股", "元件", "blue") self.assertEqual(self.database.list_watchlist(first["id"])[0]["color"], "red") self.assertEqual(self.database.list_watchlist(second["id"])[0]["color"], "blue") self.assertFalse(self.database.delete_watchlist(second["id"], "000001")) self.assertTrue(self.database.delete_watchlist(first["id"], "002141")) self.assertEqual(self.database.list_watchlist(first["id"]), []) self.assertEqual(len(self.database.list_watchlist(second["id"])), 1) def test_legacy_review_notes_are_assigned_to_first_account(self): legacy_path = Path(self.temp.name) / "legacy.db" connection = sqlite3.connect(legacy_path) try: connection.executescript( """ CREATE TABLE users ( id INTEGER PRIMARY KEY AUTOINCREMENT, username TEXT NOT NULL UNIQUE, password_salt TEXT NOT NULL, password_hash TEXT NOT NULL, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ); INSERT INTO users (username, password_salt, password_hash, created_at, updated_at) VALUES ('legacy_admin', 'salt', 'hash', '2026-01-01', '2026-01-01'); CREATE TABLE review_notes ( id INTEGER PRIMARY KEY AUTOINCREMENT, code TEXT NOT NULL DEFAULT '', stock_name TEXT NOT NULL DEFAULT '', trade_date TEXT NOT NULL, content TEXT NOT NULL DEFAULT '', plan TEXT NOT NULL DEFAULT '', created_at TEXT NOT NULL, updated_at TEXT NOT NULL ); INSERT INTO review_notes (code, stock_name, trade_date, content, plan, created_at, updated_at) VALUES ('', '', '20260721', '旧复盘', '', '2026-07-21', '2026-07-21'); CREATE TABLE watchlist ( code TEXT PRIMARY KEY, name TEXT NOT NULL, sector TEXT NOT NULL DEFAULT '', color TEXT NOT NULL DEFAULT 'red', created_at TEXT NOT NULL, updated_at TEXT NOT NULL ); INSERT INTO watchlist (code, name, sector, color, created_at, updated_at) VALUES ('002141', '贤丰控股', '元件', 'red', '2026-07-21', '2026-07-21'); """ ) connection.commit() finally: connection.close() migrated = ReviewDatabase(legacy_path) notes = migrated.list_notes(1) self.assertEqual(len(notes), 1) self.assertEqual(notes[0]["content"], "旧复盘") self.assertEqual(migrated.list_watchlist(1)[0]["code"], "002141") if __name__ == "__main__": unittest.main()