From 43b0bf293dc771e452ec046ae2eddbdbb8982da7 Mon Sep 17 00:00:00 2001 From: fengsilin Date: Wed, 17 Jun 2026 19:13:00 +0800 Subject: [PATCH] test(task-pricing): cover video usage billing flows --- test-scripts/test_video_pricing.py | 226 +++++++++++++++++++++++++++++ 1 file changed, 226 insertions(+) create mode 100644 test-scripts/test_video_pricing.py diff --git a/test-scripts/test_video_pricing.py b/test-scripts/test_video_pricing.py new file mode 100644 index 0000000..fed0812 --- /dev/null +++ b/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")