fix(scripts): push_predictions 的分页与令牌过期

两个在真实环境暴露的 bug:

1. /api/predictions/ 根本不分页——Label Studio 的 REST_FRAMEWORK 没设
   DEFAULT_PAGINATION_CLASS,PAGE_SIZE 形同虚设,它一次返回全部预测并忽略 page。
   原先按页翻永远不终止,实测翻到第 1798 页才因别的原因报错。改成只请求一次。

2. access token 默认只有 5 分钟(LS 未覆盖 SimpleJWT 默认值),而推上万条预测远超
   5 分钟。所有请求统一走 request(),碰到 401 自动用 refresh token 续期后重试一次。

另加 --replace-all:用 Data Manager 的批量操作一次删光项目里所有预测,
比逐条删同版本快得多(适合版本号变了、旧预测留着没用的情况)。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
2026-09-22 13:17:53 +08:00
co-authored by Claude Opus 5
parent 9ead8c0a6e
commit 744e5fd18a
2 changed files with 118 additions and 62 deletions
+59 -19
View File
@@ -53,9 +53,9 @@ class _FakeSession:
self.headers = {}
self.calls = []
def get(self, url, params=None, timeout=None):
def request(self, method, url, params=None, json=None, timeout=None):
self.calls.append(params)
page = params["page"]
page = (params or {}).get("page", 1)
return _FakeResponse(self.pages[page - 1] if page <= len(self.pages) else {"tasks": []})
@@ -96,36 +96,76 @@ def test_tasks_stops_on_empty_page(client):
class _TokenSession:
def __init__(self, access="ACCESS", status=200):
"""会 post 换取 access token 的假会话;request 按预设的状态码序列返回。"""
def __init__(self, access="ACCESS", status=200, statuses=None, payload=None):
self.headers = {}
self.access, self.status, self.posted = access, status, []
self.statuses = list(statuses or [])
self.payload = payload if payload is not None else []
self.requests = []
self.exchanges = 0
def post(self, url, json=None, timeout=None):
self.posted.append((url, json))
return _FakeResponse({"access": self.access}, self.status)
self.exchanges += 1
return _FakeResponse({"access": f"{self.access}{self.exchanges}"}, self.status)
def request(self, method, url, params=None, json=None, timeout=None):
self.requests.append((method, url, self.headers.get("Authorization")))
status = self.statuses.pop(0) if self.statuses else 200
return _FakeResponse(self.payload, status)
def test_legacy_token_is_sent_as_token(monkeypatch):
def _client(session, token="aaa.bbb.ccc"):
client = pp.LabelStudio.__new__(pp.LabelStudio)
client.url = "http://ls"
client.session = _TokenSession()
assert client._authorization("c657261234567890") == "Token c657261234567890"
assert client.session.posted == [] # 旧式令牌不需要换取
client.session = session
client.token = token
client._authorize()
return client
def test_jwt_personal_token_is_exchanged_for_a_bearer(monkeypatch):
def test_legacy_token_is_sent_as_token():
session = _TokenSession()
_client(session, token="c657261234567890")
assert session.headers["Authorization"] == "Token c657261234567890"
assert session.posted == [] # 旧式令牌不需要换取
def test_jwt_personal_token_is_exchanged_for_a_bearer():
"""新版默认禁用旧式 token,个人访问令牌要先换成 access token。"""
client = pp.LabelStudio.__new__(pp.LabelStudio)
client.url = "http://ls"
client.session = _TokenSession(access="AAA")
assert client._authorization("aaa.bbb.ccc") == "Bearer AAA"
assert client.session.posted[0][0] == "http://ls/api/token/refresh/"
assert client.session.posted[0][1] == {"refresh": "aaa.bbb.ccc"}
session = _TokenSession(access="AAA")
_client(session)
assert session.headers["Authorization"] == "Bearer AAA1"
assert session.posted[0] == ("http://ls/api/token/refresh/", {"refresh": "aaa.bbb.ccc"})
def test_failed_exchange_reports_clearly():
client = pp.LabelStudio.__new__(pp.LabelStudio)
client.url = "http://ls"
client.session = _TokenSession(status=401)
with pytest.raises(RuntimeError, match="access token 失败"):
client._authorization("aaa.bbb.ccc")
_client(_TokenSession(status=401))
def test_expired_access_token_is_refreshed_and_retried():
"""access token 默认只有 5 分钟,推上万条预测必然跨过期,401 要能自愈。"""
session = _TokenSession(statuses=[401, 200])
client = _client(session)
client.request("GET", "/api/tasks/")
assert session.exchanges == 2 # 初次换取 + 过期后再换一次
assert [auth for _, _, auth in session.requests] == ["Bearer ACCESS1", "Bearer ACCESS2"]
def test_predictions_endpoint_is_not_paginated():
"""回归:/api/predictions/ 一次返回全部并忽略 page,按页翻会翻不到头。"""
session = _TokenSession(payload=[
{"id": 1, "model_version": "v1"},
{"id": 2, "model_version": "v2"},
{"id": 3, "model_version": "v1"},
])
client = _client(session)
assert client.delete_predictions(39, "v1") == 2
methods = [method for method, _, _ in session.requests]
assert methods == ["GET", "DELETE", "DELETE"] # 只列一次,删两条
assert session.requests[1][1] == "http://ls/api/predictions/1/"