ModelTrainEval/tests/test_push_predictions.py

"""批量给 Label Studio 推预测：URL 还原与分页拉取。"""

import sys

import pytest

pytest.importorskip("requests")

sys.path.insert(0, "scripts")
import push_predictions as pp  # noqa: E402


def test_local_path_from_local_files_url(tmp_path):
    url = "/data/local-files/?d=boxes_shap_label_studio/images/blackdot_0489.png"
    assert pp.local_path(url, tmp_path) == (
        tmp_path / "boxes_shap_label_studio/images/blackdot_0489.png"
    )


def test_local_path_unquotes_escaped_names(tmp_path):
    """文件名里有中文或空格时 URL 是转义过的，得还原回来。"""
    url = "/data/local-files/?d=proj/images/%E5%8F%8D%E9%9D%A2%20a.png"
    assert pp.local_path(url, tmp_path) == tmp_path / "proj/images/反面 a.png"


def test_local_path_rejects_non_local_urls(tmp_path):
    with pytest.raises(ValueError, match="本地存储"):
        pp.local_path("https://example.com/a.png", tmp_path)


class _FakeResponse:
    def __init__(self, payload, status=200):
        self._payload = payload
        self.status_code = status

    @property
    def text(self):
        return str(self._payload)

    def json(self):
        return self._payload

    def raise_for_status(self):
        if self.status_code >= 400:
            raise AssertionError(f"HTTP {self.status_code}")


class _FakeSession:
    """按页返回任务，记录请求参数。"""

    def __init__(self, pages):
        self.pages = pages
        self.headers = {}
        self.calls = []

    def request(self, method, url, params=None, json=None, timeout=None):
        self.calls.append(params)
        page = (params or {}).get("page", 1)
        return _FakeResponse(self.pages[page - 1] if page <= len(self.pages) else {"tasks": []})


@pytest.fixture
def client(monkeypatch):
    monkeypatch.setattr(pp, "LabelStudio", pp.LabelStudio)
    instance = pp.LabelStudio.__new__(pp.LabelStudio)
    instance.url = "http://ls"
    return instance


def test_tasks_pagination_stops_at_total(client):
    """分页要拉完再停，不能只拿第一页。"""
    client.session = _FakeSession([
        {"tasks": [{"id": 1}, {"id": 2}], "total": 3},
        {"tasks": [{"id": 3}], "total": 3},
    ])
    assert [t["id"] for t in client.tasks(39, page_size=2)] == [1, 2, 3]


def test_tasks_handles_drf_style_pagination(client):
    """有的版本返回 results/count 而不是 tasks/total。"""
    client.session = _FakeSession([
        {"results": [{"id": 1}], "count": 2},
        {"results": [{"id": 2}], "count": 2},
    ])
    assert [t["id"] for t in client.tasks(39, page_size=1)] == [1, 2]


def test_tasks_stops_on_empty_page(client):
    """没有 total 字段时，靠空页收尾，不能无限翻。"""
    client.session = _FakeSession([{"tasks": [{"id": 1}]}, {"tasks": []}])
    assert [t["id"] for t in client.tasks(39)] == [1]
    assert len(client.session.calls) == 2


# ---------------- 两种令牌 ----------------


class _TokenSession:
    """会 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))
        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 _client(session, token="aaa.bbb.ccc"):
    client = pp.LabelStudio.__new__(pp.LabelStudio)
    client.url = "http://ls"
    client.session = session
    client.token = token
    client._authorize()
    return client


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。"""
    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():
    with pytest.raises(RuntimeError, match="access token 失败"):
        _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/"


def test_http_errors_include_the_response_body():
    """Label Studio 把真正原因写在响应体里，只抛 URL 没法排查。"""
    session = _TokenSession(statuses=[500], payload={"detail": "boom"})
    client = _client(session)
    with pytest.raises(RuntimeError, match="HTTP 500.*boom"):
        client.request("POST", "/api/projects/39/import/predictions", json=[])