fix: stream long-running website operations

2026-07-12 12:07:47 +08:00
parent c0b3b8276d
commit 40e75c0b51
6 changed files with 1248 additions and 273 deletions
@@ -327,9 +327,347 @@ def test_organize_groups_list_accounts_none_conservatively_skips(db_session, mon
    assert k3.status == "imported"


+def test_organize_groups_streaming_basic_event_sequence(db_session, monkeypatch):
+    """流式一键整理：验证基本 NDJSON 事件序列、响应头、total_items 准确性。"""
+    w = Website(
+        name="W1",
+        site_type="sub2api",
+        base_url="http://w1",
+        enabled=True,
+        auth_config_json="{}",
+        timeout_seconds=30
+    )
+    u1 = Upstream(name="U1", base_url="http://u1")
+    db_session.add_all([w, u1])
+    db_session.commit()
+    db_session.refresh(w)
+    db_session.refresh(u1)
+
+    b1 = WebsiteGroupBinding(
+        website_id=w.id,
+        target_group_id="TG1",
+        target_group_name="TG1-Group",
+        source_groups_json=json.dumps([
+            {"upstream_id": u1.id, "group_id": "G1"},
+        ]),
+        enabled=True
+    )
+    db_session.add(b1)
+
+    k1 = UpstreamGeneratedKey(
+        upstream_id=u1.id,
+        group_id="G1",
+        group_name="G1-Name",
+        key_name="Key-G1",
+        key_value="sk-g1-secret",
+        status="created",
+    )
+    db_session.add(k1)
+    db_session.commit()
+
+    created_accounts = []
+
+    class MockClient:
+        def __init__(self, **kwargs): pass
+        def __enter__(self): return self
+        def __exit__(self, *a): pass
+        def get_groups(self, *a, **kw):
+            return [{"id": "TG1", "name": "TG1-Group"}]
+        def list_accounts(self): return []
+        def extract_id(self, val): return val.get("id") if isinstance(val, dict) else str(val)
+        def create_account(self, body):
+            acc_id = f"NEW-{body['name']}"
+            new_acc = {"id": acc_id, "name": body["name"], "group_ids": body["group_ids"]}
+            created_accounts.append(new_acc)
+            return new_acc
+
+    monkeypatch.setattr("app.routers.websites.Sub2ApiWebsiteClient", MockClient)
+    monkeypatch.setattr("app.routers.websites.sync_account_priorities_for_website", lambda db, wid: [])
+
+    from fastapi.testclient import TestClient
+    from app.main import app
+    from app.database import get_db
+    from app.utils.auth import get_current_user
+
+    app.dependency_overrides[get_db] = lambda: db_session
+    app.dependency_overrides[get_current_user] = lambda: None
+
+    try:
+        client = TestClient(app)
+        resp = client.post(f"/api/websites/{w.id}/groups/organize/stream")
+        assert resp.status_code == 200
+        assert resp.headers["content-type"] == "application/x-ndjson"
+        assert resp.headers.get("cache-control") == "no-cache"
+        assert resp.headers.get("x-accel-buffering") == "no"
+
+        lines = [line if isinstance(line, str) else line.decode("utf-8") for line in resp.iter_lines() if line]
+        parsed_events = [json.loads(line) for line in lines]
+
+        # Event sequence: start → item → complete
+        assert len(parsed_events) >= 2
+        assert parsed_events[0]["event"] == "start"
+        assert parsed_events[0]["data"]["total_items"] == 1
+
+        item_events = [e for e in parsed_events if e["event"] == "item"]
+        assert len(item_events) == 1
+        assert item_events[0]["data"]["source_group_id"] == "G1"
+        assert item_events[0]["data"]["status"] == "created"
+
+        complete_event = next(e for e in parsed_events if e["event"] == "complete")
+        assert complete_event["data"]["success"] is True
+    finally:
+        app.dependency_overrides.clear()
+
+
+def test_organize_groups_streaming_409_concurrent_lock(db_session, monkeypatch):
+    """流式一键整理：同一网站并发调用返回 409 错误。"""
+    w = Website(
+        name="W1",
+        site_type="sub2api",
+        base_url="http://w1",
+        enabled=True,
+        auth_config_json="{}",
+        timeout_seconds=30
+    )
+    db_session.add(w)
+    db_session.commit()
+    db_session.refresh(w)
+
+    # Mock to make the first call run long enough to trigger concurrency
+    import time
+    original_generator = None
+
+    class MockClient:
+        def __init__(self, **kwargs): pass
+        def __enter__(self): return self
+        def __exit__(self, *a): pass
+        def get_groups(self, *a, **kw): return [{"id": "TG1", "name": "TG1-Group"}]
+        def list_accounts(self): return []
+        def extract_id(self, val): return val.get("id") if isinstance(val, dict) else str(val)
+
+    monkeypatch.setattr("app.routers.websites.Sub2ApiWebsiteClient", MockClient)
+    monkeypatch.setattr("app.routers.websites.sync_account_priorities_for_website", lambda db, wid: [])
+
+    from fastapi.testclient import TestClient
+    from app.main import app
+    from app.database import get_db
+    from app.utils.auth import get_current_user
+
+    app.dependency_overrides[get_db] = lambda: db_session
+    app.dependency_overrides[get_current_user] = lambda: None
+
+    try:
+        client = TestClient(app)
+        # First request starts the task lock
+        resp1 = client.post(f"/api/websites/{w.id}/groups/organize/stream")
+        assert resp1.status_code == 200
+
+        # Second request while first is still "running" (lock not yet released)
+        # But the first request already completed, so we need to simulate
+        # by calling the task lock manually.
+        from app.routers.websites import _try_start_task, _finish_task
+        ok, _task_token = _try_start_task(w.id, "organize")
+        assert ok
+        try:
+            resp2 = client.post(f"/api/websites/{w.id}/groups/organize/stream")
+            assert resp2.status_code == 200
+            lines = [line if isinstance(line, str) else line.decode("utf-8") for line in resp2.iter_lines() if line]
+            parsed_events = [json.loads(line) for line in lines]
+            assert len(parsed_events) == 1
+            assert parsed_events[0]["event"] == "error"
+            assert "正在执行中" in parsed_events[0]["data"]["message"]
+        finally:
+            _finish_task(w.id, "organize", _task_token)
+
+        # Task lock released: third request should succeed
+        resp3 = client.post(f"/api/websites/{w.id}/groups/organize/stream")
+        assert resp3.status_code == 200
+    finally:
+        app.dependency_overrides.clear()
+
+
+def test_organize_groups_streaming_non_sub2api_returns_error_event(db_session, monkeypatch):
+    """流式一键整理：非 sub2api 网站返回 error 事件而非 HTTP 错误。"""
+    w = Website(
+        name="W1",
+        site_type="other",
+        base_url="http://w1",
+        enabled=True,
+        auth_config_json="{}",
+        timeout_seconds=30
+    )
+    db_session.add(w)
+    db_session.commit()
+    db_session.refresh(w)
+
+    from fastapi.testclient import TestClient
+    from app.main import app
+    from app.database import get_db
+    from app.utils.auth import get_current_user
+
+    app.dependency_overrides[get_db] = lambda: db_session
+    app.dependency_overrides[get_current_user] = lambda: None
+
+    try:
+        client = TestClient(app)
+        resp = client.post(f"/api/websites/{w.id}/groups/organize/stream")
+        assert resp.status_code == 200
+        lines = [line if isinstance(line, str) else line.decode("utf-8") for line in resp.iter_lines() if line]
+        parsed_events = [json.loads(line) for line in lines]
+        assert len(parsed_events) == 1
+        assert parsed_events[0]["event"] == "error"
+    finally:
+        app.dependency_overrides.clear()
+
+
+def test_organize_groups_streaming_total_items_matches_items(db_session, monkeypatch):
+    """流式一键整理：total_items 应当等于实际产生的 item 数量。"""
+    w = Website(
+        name="W1",
+        site_type="sub2api",
+        base_url="http://w1",
+        enabled=True,
+        auth_config_json="{}",
+        timeout_seconds=30
+    )
+    u1 = Upstream(name="U1", base_url="http://u1")
+    db_session.add_all([w, u1])
+    db_session.commit()
+    db_session.refresh(w)
+    db_session.refresh(u1)
+
+    # Two bindings, each with one key
+    b1 = WebsiteGroupBinding(
+        website_id=w.id,
+        target_group_id="TG1",
+        target_group_name="TG1-Group",
+        source_groups_json=json.dumps([{"upstream_id": u1.id, "group_id": "G1"}]),
+        enabled=True
+    )
+    b2 = WebsiteGroupBinding(
+        website_id=w.id,
+        target_group_id="TG2",
+        target_group_name="TG2-Group",
+        source_groups_json=json.dumps([{"upstream_id": u1.id, "group_id": "G2"}]),
+        enabled=True
+    )
+    db_session.add_all([b1, b2])
+
+    k1 = UpstreamGeneratedKey(
+        upstream_id=u1.id, group_id="G1", group_name="G1-Name",
+        key_name="Key-G1", key_value="sk-g1", status="created",
+    )
+    k2 = UpstreamGeneratedKey(
+        upstream_id=u1.id, group_id="G2", group_name="G2-Name",
+        key_name="Key-G2", key_value="sk-g2", status="created",
+    )
+    db_session.add_all([k1, k2])
+    db_session.commit()
+
+    class MockClient:
+        def __init__(self, **kwargs): pass
+        def __enter__(self): return self
+        def __exit__(self, *a): pass
+        def get_groups(self, *a, **kw):
+            return [{"id": "TG1", "name": "TG1-Group"}, {"id": "TG2", "name": "TG2-Group"}]
+        def list_accounts(self): return []
+        def extract_id(self, val): return val.get("id") if isinstance(val, dict) else str(val)
+        def create_account(self, body):
+            return {"id": "NEW-" + body["name"], "name": body["name"], "group_ids": body["group_ids"]}
+
+    monkeypatch.setattr("app.routers.websites.Sub2ApiWebsiteClient", MockClient)
+    monkeypatch.setattr("app.routers.websites.sync_account_priorities_for_website", lambda db, wid: [])
+
+    from fastapi.testclient import TestClient
+    from app.main import app
+    from app.database import get_db
+    from app.utils.auth import get_current_user
+
+    app.dependency_overrides[get_db] = lambda: db_session
+    app.dependency_overrides[get_current_user] = lambda: None
+
+    try:
+        client = TestClient(app)
+        resp = client.post(f"/api/websites/{w.id}/groups/organize/stream")
+        assert resp.status_code == 200
+        lines = [line if isinstance(line, str) else line.decode("utf-8") for line in resp.iter_lines() if line]
+        parsed_events = [json.loads(line) for line in lines]
+
+        start_event = next(e for e in parsed_events if e["event"] == "start")
+        item_events = [e for e in parsed_events if e["event"] == "item"]
+        assert start_event["data"]["total_items"] == len(item_events)
+    finally:
+        app.dependency_overrides.clear()
+
+
+def test_organize_groups_original_json_endpoint_still_works(db_session, monkeypatch):
+    """一键整理 JSON 接口的回归测试：HTTP 状态码和错误处理仍然正确。"""
+    # Non-sub2api → HTTP 400
+    w = Website(
+        name="W1", site_type="other", base_url="http://w1",
+        enabled=True, auth_config_json="{}", timeout_seconds=30
+    )
+    db_session.add(w)
+    db_session.commit()
+    db_session.refresh(w)
+
+    from fastapi.testclient import TestClient
+    from app.main import app
+    from app.database import get_db
+    from app.utils.auth import get_current_user
+
+    app.dependency_overrides[get_db] = lambda: db_session
+    app.dependency_overrides[get_current_user] = lambda: None
+
+    try:
+        client = TestClient(app)
+        resp = client.post(f"/api/websites/{w.id}/groups/organize")
+        assert resp.status_code == 400
+        assert "只支持" in resp.json()["detail"] or "only sub2api" in resp.json()["detail"]
+    finally:
+        app.dependency_overrides.clear()
+
+
+def test_organize_groups_json_409_concurrent_lock(db_session, monkeypatch):
+    """一键整理 JSON 接口：并发调用返回 HTTP 409。"""
+    w = Website(
+        name="W1", site_type="sub2api", base_url="http://w1",
+        enabled=True, auth_config_json="{}", timeout_seconds=30,
+    )
+    db_session.add(w)
+    db_session.commit()
+    db_session.refresh(w)
+
+    from fastapi.testclient import TestClient
+    from app.main import app
+    from app.database import get_db
+    from app.utils.auth import get_current_user
+
+    app.dependency_overrides[get_db] = lambda: db_session
+    app.dependency_overrides[get_current_user] = lambda: None
+
+    try:
+        client = TestClient(app)
+        # Hold lock manually to simulate concurrent request
+        from app.routers.websites import _try_start_task, _finish_task
+        ok, _task_token = _try_start_task(w.id, "organize")
+        assert ok
+        try:
+            resp = client.post(f"/api/websites/{w.id}/groups/organize")
+            assert resp.status_code == 409
+            assert "正在执行中" in resp.json()["detail"]
+        finally:
+            _finish_task(w.id, "organize", _task_token)
+
+        # Lock released: second request should succeed
+        resp2 = client.post(f"/api/websites/{w.id}/groups/organize")
+        assert resp2.status_code == 200
+    finally:
+        app.dependency_overrides.clear()
+
+
 def test_organize_groups_force_alignment(db_session, monkeypatch):
    """测试强对齐场景：
-    1. 已存在账号在旧 SmartUp 分组，整理后移除旧分组并加入新分组。
    2. 已存在账号同时有非 SmartUp 分组，整理后保留非 SmartUp 分组。
    3. 已存在账号已经完全一致，整理不重复调用更新接口。
    4. 多目标分组、多上游来源时，每个账号按自己的当前绑定关系对齐。