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:
co-authored by
Claude Opus 5
parent
9ead8c0a6e
commit
744e5fd18a
@@ -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/"
|
||||
|
||||
Reference in New Issue
Block a user