fix: 防止分页漏读和上游异常快照误标 orphaned

2026-07-02 14:21:09 +08:00
parent 00af34c487
commit e2129d3796
4 changed files with 349 additions and 9 deletions
@@ -10,6 +10,7 @@ from sqlalchemy.pool import StaticPool
 from app.database import Base
 from app.models.upstream import Upstream
 from app.models.website import Website  # noqa: F401 — registers table for FK refs
+from app.models.snapshot import UpstreamRateSnapshot  # noqa: F401 — registers table for FK refs
 from app.models.upstream_key import UpstreamGeneratedKey


@@ -1834,3 +1835,200 @@ def test_ensure_group_key_rollback_prevents_leak(db_session):
    assert vip_db.key_value == "sk-vip-local"


+def test_upstream_key_reconciliation_robustness(db_session, monkeypatch):
+    """测试对账逻辑的健壮性：单次快照缺失不标记、两次连续快照缺失标记、自恢复、拉取失败时不误清理/误恢复。"""
+    from app.services import scheduler as sched_mod
+    from app.models.upstream_key import UpstreamGeneratedKey
+    from app.models.snapshot import UpstreamRateSnapshot
+    from app.services.upstream_client import UpstreamClient
+    from app.services import website_sync
+    import json
+    from datetime import datetime, timezone
+
+    # 1. 准备基础数据
+    website = Website(
+        name="Target",
+        site_type="sub2api",
+        base_url="http://target.local",
+        api_prefix="/api/v1/admin",
+        auth_type="api_key",
+        auth_config_json="{}",
+    )
+    upstream = Upstream(
+        name="TestUpstream",
+        base_url="http://upstream.local",
+        api_prefix="/api/v1",
+        auth_type="bearer",
+        auth_config_json="{}",
+    )
+    db_session.add_all([website, upstream])
+    db_session.commit()
+    db_session.refresh(website)
+    db_session.refresh(upstream)
+
+    # 验证 UpstreamRateSnapshot 表是否成功创建并可正常读写
+    snap = UpstreamRateSnapshot(
+        upstream_id=upstream.id,
+        snapshot_json="{}",
+        captured_at=datetime.now(timezone.utc)
+    )
+    db_session.add(snap)
+    db_session.commit()
+    assert db_session.query(UpstreamRateSnapshot).count() == 1
+
+    # 插入一个初始的已导入 key，目前状态是 imported
+    key_row = UpstreamGeneratedKey(
+        upstream_id=upstream.id,
+        group_id="plus_group",
+        group_name="Plus智能分组",
+        key_name="SmartUp-TestUpstream-plus_group",
+        key_value="sk-plus-value",
+        managed_prefix="SmartUp",
+        key_id="key-12345",
+        status="imported",
+        imported_website_id=website.id,
+        imported_account_id="acc-999",
+        imported_target_group_id="target-3",
+    )
+    db_session.add(key_row)
+    db_session.commit()
+    db_session.refresh(key_row)
+
+    # 准备 Mock UpstreamClient
+    mock_keys = [{"id": "key-12345", "name": "SmartUp-TestUpstream-plus_group", "group_id": "plus_group"}]
+    monkeypatch.setattr(UpstreamClient, "list_api_keys", lambda self, **kw: mock_keys)
+    monkeypatch.setattr(UpstreamClient, "login", lambda self: None)
+    monkeypatch.setattr(UpstreamClient, "close", lambda self: None)
+    monkeypatch.setattr(UpstreamClient, "__enter__", lambda self: self)
+    monkeypatch.setattr(UpstreamClient, "__exit__", lambda self, *a: None)
+    monkeypatch.setattr(sched_mod, "SessionLocal", lambda: db_session)
+    original_close = db_session.close
+    monkeypatch.setattr(db_session, "close", lambda: None)
+
+    try:
+        # === 场景 A: 单次异常快照缺少来源分组 ===
+        # 运行对账，传入前一次快照存在该分组，当前快照不含该分组
+        active_group_ids_1 = {"other_group"}
+        prev_active_group_ids_1 = {"plus_group", "other_group"}
+        remote_key_ids_1 = {"key-12345"}
+        captured_at_1 = datetime.now(timezone.utc)
+
+        website_sync.reconcile_upstream_keys(
+            db_session,
+            upstream.id,
+            active_group_ids_1,
+            remote_key_ids_1,
+            captured_at_1,
+            prev_active_group_ids=prev_active_group_ids_1
+        )
+        db_session.commit()
+
+        # 验证：单次异常快照缺少该分组，key 不应被标记为 orphaned
+        db_session.refresh(key_row)
+        assert key_row.status == "imported", "单次缺失分组不应该将 key 标为 orphaned"
+        assert key_row.error is None
+
+        # === 场景 B: 两次连续快照都缺少来源分组 ===
+        # 运行对账，传入前一个和当前快照都不含该分组
+        active_group_ids_2 = {"other_group"}
+        prev_active_group_ids_2 = {"other_group"}
+        remote_key_ids_2 = {"key-12345"}
+        captured_at_2 = datetime.now(timezone.utc)
+
+        website_sync.reconcile_upstream_keys(
+            db_session,
+            upstream.id,
+            active_group_ids_2,
+            remote_key_ids_2,
+            captured_at_2,
+            prev_active_group_ids=prev_active_group_ids_2
+        )
+        db_session.commit()
+
+        # 验证：两次连续快照都缺少该分组，应标记为 orphaned
+        db_session.refresh(key_row)
+        assert key_row.status == "orphaned", "连续两次缺失分组应该将 key 标为 orphaned"
+        assert key_row.error == "来源分组已不存在"
+
+        # === 场景 C: 自恢复逻辑 ===
+        # 运行对账，传入当前快照分组已恢复，且远端 key 存在
+        active_group_ids_3 = {"plus_group", "other_group"}
+        remote_key_ids_3 = {"key-12345"}
+        captured_at_3 = datetime.now(timezone.utc)
+
+        website_sync.reconcile_upstream_keys(
+            db_session,
+            upstream.id,
+            active_group_ids_3,
+            remote_key_ids_3,
+            captured_at_3
+        )
+        db_session.commit()
+
+        # 验证：分组恢复且远端 key id 存在时，应恢复为 imported 并清空错误
+        db_session.refresh(key_row)
+        assert key_row.status == "imported", "恢复分组且远端 key 存在时应该恢复为 imported"
+        assert key_row.error is None
+
+        # === 场景 D: 远端 key 列表拉取失败时，不做 key id 级清理或恢复 ===
+        key_row.status = "orphaned"
+        key_row.error = "临时标记"
+        db_session.commit()
+
+        active_group_ids_4 = {"plus_group", "other_group"}
+        remote_key_ids_4 = None
+        captured_at_4 = datetime.now(timezone.utc)
+
+        website_sync.reconcile_upstream_keys(
+            db_session,
+            upstream.id,
+            active_group_ids_4,
+            remote_key_ids_4,
+            captured_at_4
+        )
+        db_session.commit()
+
+        db_session.refresh(key_row)
+        assert key_row.status == "orphaned"
+        assert key_row.error == "临时标记"
+
+        # === 场景 E: 重复旧 key 级防误恢复 ===
+        # 创建一个重复旧 key (orphaned，error="重复旧 Key...", imported_account_id=None)
+        dup_key_row = UpstreamGeneratedKey(
+            upstream_id=upstream.id,
+            group_id="plus_group",
+            group_name="Plus智能分组",
+            key_name="SmartUp-TestUpstream-plus_group-dup",
+            key_value="sk-dup-value",
+            managed_prefix="SmartUp",
+            key_id="key-dup-999",
+            status="orphaned",
+            error="重复旧 Key；已由 key #73 对应当前分组导入账号",
+            imported_website_id=None,
+            imported_account_id=None,
+            imported_target_group_id=None,
+        )
+        db_session.add(dup_key_row)
+        db_session.commit()
+
+        # 运行对账，传入分组恢复且远端存在 key-dup-999
+        active_group_ids_5 = {"plus_group", "other_group"}
+        remote_key_ids_5 = {"key-12345", "key-dup-999"}
+        captured_at_5 = datetime.now(timezone.utc)
+
+        website_sync.reconcile_upstream_keys(
+            db_session,
+            upstream.id,
+            active_group_ids_5,
+            remote_key_ids_5,
+            captured_at_5
+        )
+        db_session.commit()
+
+        # 验证：即使分组和 key_id 存在，因为它是重复旧 key 且 imported_account_id=None，也必须保持 orphaned 并且 error 不变
+        db_session.refresh(dup_key_row)
+        assert dup_key_row.status == "orphaned", "重复旧 key 即使恢复也不应被误标记为 imported"
+        assert dup_key_row.error == "重复旧 Key；已由 key #73 对应当前分组导入账号"
+
+    finally:
+        monkeypatch.setattr(db_session, "close", original_close)