xiaoxia-saas/scripts/smoke_test.py

#!/usr/bin/env python3
"""
冒烟测试脚本 - 部署后自动验证核心端点可用性
用法: python3 smoke_test.py <API_BASE_URL> [--email EMAIL] [--password PASSWORD] [--json]
示例: python3 smoke_test.py https://saas-api.xiaoxiajianji.com --email test@example.com --password test123 --json
"""

import argparse
import json
import logging
import ssl
import sys
import time
import urllib.error
import urllib.request

logger = logging.getLogger(__name__)

CORE_ENDPOINTS = [
    {
        "name": "upload/direct/prepare",
        "method": "POST",
        "path": "/api/v1/upload/direct/prepare",
        "body": {"project_id": "smoke-test", "file_name": "test.mp4", "file_size": 1024, "content_type": "video/mp4"},
        "expect": [200, 401, 422],
    },
    {
        "name": "upload/chunk/init",
        "method": "POST",
        "path": "/api/v1/upload/chunk/init",
        "body": {"project_id": "smoke-test", "file_name": "test.mp4", "file_size": 1024000, "total_chunks": 2},
        "expect": [200, 401, 422],
    },
    {"name": "dashboard/overview", "method": "GET", "path": "/api/v1/dashboard/overview", "expect": [200, 401]},
    {"name": "assets", "method": "GET", "path": "/api/v1/assets?library_id=smoke-test", "expect": [200, 401]},
    {
        "name": "generation/tasks",
        "method": "POST",
        "path": "/api/v1/generation/tasks",
        "body": {},
        "expect": [200, 401, 422],
    },
]


def make_request(base_url, endpoint, token=None):
    url = f"{base_url}{endpoint['path']}"
    headers = {"Content-Type": "application/json"}
    if token:
        headers["Authorization"] = f"Bearer {token}"

    data = json.dumps(endpoint.get("body", {})).encode() if endpoint.get("body") is not None else None
    req = urllib.request.Request(url, data=data, headers=headers, method=endpoint["method"])

    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE

    try:
        start = time.time()
        resp = urllib.request.urlopen(req, timeout=15, context=ctx)
        elapsed = round((time.time() - start) * 1000)
        body = resp.read().decode()
        return {"status": resp.status, "elapsed_ms": elapsed, "body": body[:200], "error": None}
    except urllib.error.HTTPError as e:
        elapsed = round((time.time() - start) * 1000)
        body = ""
        try:
            body = e.read().decode()[:200]
        except:
            logger.warning(f"Operation failed in scripts/smoke_test.py: {e}", exc_info=True)
        return {"status": e.code, "elapsed_ms": elapsed, "body": body, "error": None}
    except Exception as e:
        return {"status": 0, "elapsed_ms": 0, "body": "", "error": str(e)}


def login(base_url, email, password):
    url = f"{base_url}/api/v1/auth/login"
    data = json.dumps({"email": email, "password": password}).encode()
    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}, method="POST")
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    try:
        resp = urllib.request.urlopen(req, timeout=10, context=ctx)
        body = json.loads(resp.read().decode())
        return body.get("token") or body.get("data", {}).get("token") or body.get("access_token")
    except:
        return None


def run_smoke_test(base_url, email=None, password=None, output_json=False):
    base_url = base_url.rstrip("/")
    token = None

    if email and password:
        token = login(base_url, email, password)
        if not output_json:
            print(f"{'✅ 登录成功' if token else '⚠️ 登录失败，将以未认证模式测试'}")

    results = []
    all_passed = True

    for ep in CORE_ENDPOINTS:
        result = make_request(base_url, ep, token)
        passed = result["status"] in ep["expect"] and result["error"] is None
        is_5xx = 500 <= result["status"] < 600
        if is_5xx:
            passed = False
            all_passed = False

        results.append(
            {
                "name": ep["name"],
                "path": ep["path"],
                "status": result["status"],
                "elapsed_ms": result["elapsed_ms"],
                "passed": passed,
                "error": result["error"],
                "is_5xx": is_5xx,
            }
        )

        if not output_json:
            icon = "✅" if passed else "❌"
            print(f"  {icon} {ep['name']}: {result['status']} ({result['elapsed_ms']}ms)")

    if output_json:
        print(json.dumps({"success": all_passed, "results": results, "base_url": base_url}, indent=2))

    return 0 if all_passed else 1


def main():
    parser = argparse.ArgumentParser(description="冒烟测试 - 部署后核心端点验证")
    parser.add_argument("base_url", help="API 基础地址，如 https://saas-api.xiaoxiajianji.com")
    parser.add_argument("--email", help="登录邮箱")
    parser.add_argument("--password", help="登录密码")
    parser.add_argument("--json", action="store_true", help="JSON 格式输出")
    args = parser.parse_args()

    if not args.json:
        print(f"\n🔍 冒烟测试: {args.base_url}")
        print("-" * 50)

    exit_code = run_smoke_test(args.base_url, args.email, args.password, args.json)

    if not args.json:
        print("-" * 50)
        print(f"{'✅ 全部通过' if exit_code == 0 else '❌ 存在失败端点'}\n")

    sys.exit(exit_code)


if __name__ == "__main__":
    main()