diff --git a/backend/app/routers/websites.py b/backend/app/routers/websites.py index b000024..cad33dc 100644 --- a/backend/app/routers/websites.py +++ b/backend/app/routers/websites.py @@ -1123,156 +1123,161 @@ def set_website_accounts_concurrency( if website.site_type != "sub2api": raise HTTPException(400, "only sub2api site supports account concurrency settings") - c = _client(website) + with _client(website) as c: + try: + remote_accounts = c.list_accounts() + except Exception as e: + return SetConcurrencyResponse( + success=False, + message=f"拉取远端账号列表失败: {e}", + items=[] + ) - try: - remote_accounts = c.list_accounts() - except Exception as e: - return SetConcurrencyResponse( - success=False, - message=f"拉取远端账号列表失败: {e}", - items=[] - ) + if remote_accounts is None: + return SetConcurrencyResponse( + success=False, + message="拉取远端账号列表失败，无法设置并发", + items=[] + ) - remote_map = {} - if remote_accounts: + remote_map = {} for acc in remote_accounts: aid = c.extract_id(acc) if aid: remote_map[aid] = acc - keys = db.query(UpstreamGeneratedKey).filter( - UpstreamGeneratedKey.imported_website_id == wid, - UpstreamGeneratedKey.imported_account_id.isnot(None) - ).all() + keys = db.query(UpstreamGeneratedKey).filter( + UpstreamGeneratedKey.imported_website_id == wid, + UpstreamGeneratedKey.imported_account_id.isnot(None) + ).all() - candidates = {} - for key in keys: - aid = key.imported_account_id - if aid not in candidates: - candidates[aid] = { - "account_id": aid, - "db_key": key - } + candidates = {} + for key in keys: + aid = key.imported_account_id + if aid not in candidates: + candidates[aid] = { + "account_id": aid, + "db_key": key + } - if not candidates: - return SetConcurrencyResponse( - success=True, - message="没有找到 SmartUp 导入的有效账号", - items=[] - ) + if not candidates: + return SetConcurrencyResponse( + success=True, + message="没有找到 SmartUp 导入的有效账号", + items=[] + ) - bulk_candidates = [] - skipped_items = [] + bulk_candidates = [] + skipped_items = [] - for aid, cand in candidates.items(): - remote_acc = remote_map.get(aid) - if not remote_acc: - skipped_items.append(SetConcurrencyItem( - account_id=aid, - account_name=None, - old_concurrency=None, - target_concurrency=body.concurrency, - status="skipped", - message="账号在远端已被删除或不存在" - )) - continue + for aid, cand in candidates.items(): + remote_acc = remote_map.get(aid) + if not remote_acc: + skipped_items.append(SetConcurrencyItem( + account_id=aid, + account_name=None, + old_concurrency=None, + target_concurrency=body.concurrency, + status="skipped", + message="账号在远端已被删除或不存在" + )) + continue - acc_name = remote_acc.get("name") - old_con = remote_acc.get("concurrency") + acc_name = remote_acc.get("name") + old_con = remote_acc.get("concurrency") + + try: + int(aid) + bulk_candidates.append({ + "account_id": aid, + "account_name": acc_name, + "old_concurrency": old_con + }) + except ValueError: + skipped_items.append(SetConcurrencyItem( + account_id=aid, + account_name=acc_name, + old_concurrency=old_con, + target_concurrency=body.concurrency, + status="skipped", + message="账号 ID 无法转换为数字，跳过批量更新" + )) + + items = [] + items.extend(skipped_items) + + if not bulk_candidates: + success = all(item.status == "skipped" for item in items) + return SetConcurrencyResponse( + success=success, + message=f"无可执行批量更新 of 账号。已跳过 {len(skipped_items)} 个", + items=items + ) + + bulk_ids = [item["account_id"] for item in bulk_candidates] + bulk_success = False + bulk_error_msg = "" try: - int(aid) - bulk_candidates.append({ - "account_id": aid, - "account_name": acc_name, - "old_concurrency": old_con - }) - except ValueError: - skipped_items.append(SetConcurrencyItem( - account_id=aid, - account_name=acc_name, - old_concurrency=old_con, - target_concurrency=body.concurrency, - status="skipped", - message="账号 ID 无法转换为数字，跳过批量更新" - )) + c.bulk_update_accounts(bulk_ids, {"concurrency": body.concurrency}) + bulk_success = True + except Exception as e: + bulk_error_msg = str(e) + logger.warning("bulk_update_accounts failed, falling back to individual updates: %s", e) - items = [] - items.extend(skipped_items) - - if not bulk_candidates: - success = all(item.status == "skipped" for item in items) - return SetConcurrencyResponse( - success=success, - message=f"无可执行批量更新的账号。已跳过 {len(skipped_items)} 个", - items=items - ) - - bulk_ids = [item["account_id"] for item in bulk_candidates] - bulk_success = False - bulk_error_msg = "" - - try: - c.bulk_update_accounts(bulk_ids, {"concurrency": body.concurrency}) - bulk_success = True - except Exception as e: - bulk_error_msg = str(e) - logger.warning("bulk_update_accounts failed, falling back to individual updates: %s", e) - - if bulk_success: - for item in bulk_candidates: - items.append(SetConcurrencyItem( - account_id=item["account_id"], - account_name=item["account_name"], - old_concurrency=item["old_concurrency"], - target_concurrency=body.concurrency, - status="success", - message="成功" - )) - else: - for item in bulk_candidates: - aid = item["account_id"] - try: - c.update_account(aid, {"concurrency": body.concurrency}) + if bulk_success: + for item in bulk_candidates: items.append(SetConcurrencyItem( - account_id=aid, + account_id=item["account_id"], account_name=item["account_name"], old_concurrency=item["old_concurrency"], target_concurrency=body.concurrency, status="success", - message="成功 (逐个更新)" - )) - except Exception as e: - items.append(SetConcurrencyItem( - account_id=aid, - account_name=item["account_name"], - old_concurrency=item["old_concurrency"], - target_concurrency=body.concurrency, - status="failed", - message=f"更新失败: {e}" + message="成功" )) + else: + for item in bulk_candidates: + aid = item["account_id"] + try: + c.update_account(aid, {"concurrency": body.concurrency}) + items.append(SetConcurrencyItem( + account_id=aid, + account_name=item["account_name"], + old_concurrency=item["old_concurrency"], + target_concurrency=body.concurrency, + status="success", + message="成功 (逐个更新)" + )) + except Exception as e: + items.append(SetConcurrencyItem( + account_id=aid, + account_name=item["account_name"], + old_concurrency=item["old_concurrency"], + target_concurrency=body.concurrency, + status="failed", + message=f"更新失败: {e}" + )) - success_count = sum(1 for item in items if item.status == "success") - failed_count = sum(1 for item in items if item.status == "failed") - skip_count = sum(1 for item in items if item.status == "skipped") + success_count = sum(1 for item in items if item.status == "success") + failed_count = sum(1 for item in items if item.status == "failed") + skip_count = sum(1 for item in items if item.status == "skipped") - success = (failed_count == 0) - msg_parts = [] - if success_count: - msg_parts.append(f"成功 {success_count} 个") - if failed_count: - msg_parts.append(f"失败 {failed_count} 个") - if skip_count: - msg_parts.append(f"跳过 {skip_count} 个") + success = (failed_count == 0) + msg_parts = [] + if success_count: + msg_parts.append(f"成功 {success_count} 个") + if failed_count: + msg_parts.append(f"失败 {failed_count} 个") + if skip_count: + msg_parts.append(f"跳过 {skip_count} 个") - message = "设置账号并发执行完毕：" + "，".join(msg_parts) + message = "设置账号并发执行完毕：" + "，".join(msg_parts) - return SetConcurrencyResponse( - success=success, - message=message, - items=items - ) + return SetConcurrencyResponse( + success=success, + message=message, + items=items + ) @router.post("/api/websites/{wid}/groups/organize", response_model=OrganizeGroupsResponse) diff --git a/backend/test_account_concurrency.py b/backend/test_account_concurrency.py index 7c3c275..dccfa09 100644 --- a/backend/test_account_concurrency.py +++ b/backend/test_account_concurrency.py @@ -98,8 +98,17 @@ def test_set_website_accounts_concurrency_workflow(db_session, monkeypatch): bulk_calls = [] update_calls = [] list_calls_count = 0 + closed_count = 0 class MockClient: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + nonlocal closed_count + closed_count += 1 + return False + def list_accounts(self): nonlocal list_calls_count list_calls_count += 1 @@ -146,6 +155,8 @@ def test_set_website_accounts_concurrency_workflow(db_session, monkeypatch): assert len(bulk_calls) == 1 assert bulk_calls[0] == (["1001"], {"concurrency": 200}) assert len(update_calls) == 0 + # 验证是否调用了 __exit__ 关闭连接 + assert closed_count == 1 def test_set_website_accounts_concurrency_fallback(db_session, monkeypatch): @@ -179,8 +190,17 @@ def test_set_website_accounts_concurrency_fallback(db_session, monkeypatch): bulk_calls = [] update_calls = [] + closed_count = 0 class MockClientError: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + nonlocal closed_count + closed_count += 1 + return False + def list_accounts(self): return [{"id": 1001, "name": "acc-1001", "concurrency": 10}] @@ -210,3 +230,42 @@ def test_set_website_accounts_concurrency_fallback(db_session, monkeypatch): assert len(items) == 1 assert items[0].status == "success" assert "逐个更新" in items[0].message + assert closed_count == 1 + + +def test_set_website_accounts_concurrency_list_accounts_none(db_session, monkeypatch): + w = Website( + id=1, + name="Sub2Api site", + site_type="sub2api", + base_url="http://sub2api", + api_prefix="api/v1", + auth_type="bearer", + auth_config_json='{"token": "tok1"}' + ) + db_session.add(w) + db_session.commit() + + closed_count = 0 + + class MockClientNone: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + nonlocal closed_count + closed_count += 1 + return False + + def list_accounts(self): + return None + + monkeypatch.setattr("app.routers.websites._client", lambda row: MockClientNone()) + + req = SetConcurrencyRequest(concurrency=150) + res = set_website_accounts_concurrency(wid=w.id, body=req, db=db_session) + + assert res.success is False + assert "拉取远端账号列表失败" in res.message + assert len(res.items) == 0 + assert closed_count == 1