test(task-pricing): cover video usage billing flows

test-scripts/test_video_pricing.py

@@ -0,0 +1,226 @@
+#!/usr/bin/env python3
+"""
+Smoke test for video pricing table billing.
+
+Required environment:
+  NEW_API_BASE_URL=http://127.0.0.1:3000
+  NEW_API_ADMIN_USERNAME=root
+  NEW_API_ADMIN_PASSWORD=123456
+  NEW_API_TOKEN=sk-...
+
+Optional model overrides:
+  VIDEO_PRICING_HAILUO_MODEL=hailuo-video
+  VIDEO_PRICING_SEEDANCE_MODEL=seedance-2
+
+The server must already have usable channels for the selected models. For the
+per_1m_tokens case, point the Seedance/Doubao channel at a mock upstream that
+returns a completed task with usage.total_tokens.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+import time
+import urllib.error
+import urllib.parse
+import urllib.request
+from http.cookiejar import CookieJar
+
+
+BASE_URL = os.getenv("NEW_API_BASE_URL", "http://127.0.0.1:3000").rstrip("/")
+ADMIN_USERNAME = os.getenv("NEW_API_ADMIN_USERNAME", "root")
+ADMIN_PASSWORD = os.getenv("NEW_API_ADMIN_PASSWORD", "123456")
+API_TOKEN = os.getenv("NEW_API_TOKEN", "")
+HAILUO_MODEL = os.getenv("VIDEO_PRICING_HAILUO_MODEL", "hailuo-video")
+SEEDANCE_MODEL = os.getenv("VIDEO_PRICING_SEEDANCE_MODEL", "seedance-2")
+
+QUOTA_PER_UNIT = int(os.getenv("VIDEO_PRICING_QUOTA_PER_UNIT", "500000"))
+GROUP_RATIO = float(os.getenv("VIDEO_PRICING_GROUP_RATIO", "1"))
+
+
+class Client:
+    def __init__(self) -> None:
+        self.cookies = CookieJar()
+        self.opener = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(self.cookies))
+
+    def request(self, method: str, path: str, body=None, token: str | None = None, expected: int | None = None):
+        url = f"{BASE_URL}{path}"
+        data = None
+        headers = {"Content-Type": "application/json"}
+        if body is not None:
+            data = json.dumps(body).encode("utf-8")
+        if token:
+            headers["Authorization"] = f"Bearer {token}"
+        req = urllib.request.Request(url, data=data, headers=headers, method=method)
+        try:
+            with self.opener.open(req, timeout=30) as resp:
+                payload = read_json(resp)
+                if expected is not None and resp.status != expected:
+                    raise AssertionError(f"{method} {path}: expected {expected}, got {resp.status}: {payload}")
+                return resp.status, payload
+        except urllib.error.HTTPError as err:
+            payload = read_json(err)
+            if expected is not None and err.code == expected:
+                return err.code, payload
+            raise AssertionError(f"{method} {path}: HTTP {err.code}: {payload}") from err
+
+
+def read_json(resp):
+    raw = resp.read().decode("utf-8", errors="replace")
+    if not raw:
+        return None
+    try:
+        return json.loads(raw)
+    except json.JSONDecodeError:
+        return raw
+
+
+def assert_success(payload, context: str):
+    if not isinstance(payload, dict) or payload.get("success") is not True:
+        raise AssertionError(f"{context} failed: {payload}")
+    return payload.get("data")
+
+
+def usage(client: Client) -> dict:
+    _, payload = client.request("GET", "/api/usage/token/", token=API_TOKEN, expected=200)
+    if not payload.get("code"):
+        raise AssertionError(f"usage failed: {payload}")
+    return payload["data"]
+
+
+def used_quota(client: Client) -> int:
+    return int(usage(client)["total_used"])
+
+
+def login(client: Client) -> None:
+    _, payload = client.request(
+        "POST",
+        "/api/user/login",
+        {"username": ADMIN_USERNAME, "password": ADMIN_PASSWORD},
+        expected=200,
+    )
+    assert_success(payload, "admin login")
+
+
+def put_model_pricing(client: Client, model: str, config: dict) -> None:
+    _, payload = client.request("PUT", f"/api/option/model_pricing/{urllib.parse.quote(model, safe='')}", config, expected=200)
+    assert_success(payload, f"put pricing {model}")
+
+
+def delete_model_pricing(client: Client, model: str) -> None:
+    _, payload = client.request("DELETE", f"/api/option/model_pricing/{urllib.parse.quote(model, safe='')}", expected=200)
+    assert_success(payload, f"delete pricing {model}")
+
+
+def submit_video(client: Client, body: dict, expected: int = 200):
+    _, payload = client.request("POST", "/v1/videos", body, token=API_TOKEN, expected=expected)
+    return payload
+
+
+def fetch_video(client: Client, task_id: str):
+    _, payload = client.request("GET", f"/v1/videos/{urllib.parse.quote(task_id, safe='')}", token=API_TOKEN, expected=200)
+    return payload
+
+
+def hailuo_config() -> dict:
+    return {
+        "schema_version": 1,
+        "scope": "model",
+        "billing_unit": "per_call",
+        "preconsume_strategy": "exact",
+        "dimensions": [
+            {"key": "resolution", "source": "request.resolution", "type": "string", "default": None},
+            {"key": "duration", "source": "request.duration", "type": "number", "default": None},
+        ],
+        "table": [
+            {"resolution": "768P", "duration": 10, "price": 0.04, "source": "manual"},
+        ],
+        "fallback": {"strategy": "reject"},
+    }
+
+
+def seedance_config() -> dict:
+    return {
+        "schema_version": 1,
+        "scope": "model",
+        "billing_unit": "per_1m_tokens",
+        "preconsume_strategy": "minimum",
+        "dimensions": [
+            {"key": "resolution", "source": "request.resolution", "type": "string", "default": None},
+        ],
+        "table": [
+            {"resolution": "720P", "price": 0.5, "source": "manual"},
+        ],
+        "fallback": {"strategy": "reject"},
+    }
+
+
+def expected_quota(price_usd: float, tokens: int | None = None) -> int:
+    units = tokens if tokens is not None else 1
+    divisor = 1_000_000 if tokens is not None else 1
+    quota = int(float(units) * price_usd / divisor * QUOTA_PER_UNIT * GROUP_RATIO)
+    if tokens and quota <= 0 and price_usd > 0 and GROUP_RATIO > 0:
+        return 1
+    return quota
+
+
+def run() -> None:
+    if not API_TOKEN:
+        raise RuntimeError("NEW_API_TOKEN is required")
+
+    client = Client()
+    login(client)
+
+    try:
+        put_model_pricing(client, HAILUO_MODEL, hailuo_config())
+        before = used_quota(client)
+        payload = submit_video(
+            client,
+            {"model": HAILUO_MODEL, "prompt": "pricing smoke", "resolution": "768P", "duration": 10},
+        )
+        after = used_quota(client)
+        assert after - before == expected_quota(0.04), f"Hailuo quota mismatch: before={before}, after={after}, response={payload}"
+
+        before_reject = used_quota(client)
+        rejected = submit_video(
+            client,
+            {"model": HAILUO_MODEL, "prompt": "pricing smoke", "resolution": "1080P", "duration": 10},
+            expected=400,
+        )
+        after_reject = used_quota(client)
+        assert after_reject == before_reject, f"reject changed quota: {before_reject} -> {after_reject}, response={rejected}"
+
+        put_model_pricing(client, SEEDANCE_MODEL, seedance_config())
+        before_usage = used_quota(client)
+        seedance_submit = submit_video(
+            client,
+            {"model": SEEDANCE_MODEL, "prompt": "pricing smoke", "resolution": "720P", "duration": 5},
+        )
+        task_id = seedance_submit.get("id") or seedance_submit.get("task_id") or seedance_submit.get("data", {}).get("id")
+        if not task_id:
+            raise AssertionError(f"Seedance submit did not return task id: {seedance_submit}")
+
+        total_tokens = int(os.getenv("VIDEO_PRICING_EXPECTED_TOTAL_TOKENS", "1000"))
+        for _ in range(12):
+            fetched = fetch_video(client, task_id)
+            status = fetched.get("status") or fetched.get("data", {}).get("status")
+            if status in {"succeeded", "completed", "success"}:
+                break
+            time.sleep(1)
+        after_usage = used_quota(client)
+        delta = after_usage - before_usage
+        assert delta == expected_quota(0.5, total_tokens), f"Seedance usage quota mismatch: delta={delta}, response={fetched}"
+    finally:
+        delete_model_pricing(client, HAILUO_MODEL)
+        delete_model_pricing(client, SEEDANCE_MODEL)
+
+
+if __name__ == "__main__":
+    try:
+        run()
+    except Exception as exc:
+        print(f"FAIL: {exc}", file=sys.stderr)
+        sys.exit(1)
+    print("PASS: video pricing smoke completed")