smoke tester

test_serve_smoke.py

#!/usr/bin/env python3
"""
Smoke tests for a running ``transformers serve`` instance.

Usage:
    python test_serve_smoke.py [BASE_URL]
    python test_serve_smoke.py http://localhost:8000
"""

import json
import sys

import httpx

BASE = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:8000"

PASS = 0
FAIL = 0


def green(msg):
    print(f"\033[32m  PASS: {msg}\033[0m")


def red(msg):
    print(f"\033[31m  FAIL: {msg}\033[0m")


def assert_contains(label, haystack, needle):
    global PASS, FAIL
    if needle in haystack:
        green(label)
        PASS += 1
    else:
        red(f"{label} — expected {needle!r}")
        FAIL += 1


def assert_true(label, condition):
    global PASS, FAIL
    if condition:
        green(label)
        PASS += 1
    else:
        red(label)
        FAIL += 1


DEFAULT_MODEL = "HuggingFaceTB/SmolLM2-135M-Instruct"


def pick_model(client):
    """Use the default model if available, otherwise pick the first instruct model."""
    resp = client.get(f"{BASE}/v1/models")
    resp.raise_for_status()
    models = [m["id"] for m in resp.json()["data"]]
    if DEFAULT_MODEL in models:
        return DEFAULT_MODEL
    instruct = [m for m in models if "instruct" in m.lower()]
    return instruct[0] if instruct else models[0]


def chat_request(model, stream=False, max_tokens=20, content="Say hi in 3 words."):
    return {
        "model": model,
        "messages": [{"role": "user", "content": content}],
        "stream": stream,
        "max_tokens": max_tokens,
    }


def test_health(client):
    print("=== Health check ===")
    resp = client.get(f"{BASE}/health")
    assert_true("status 200", resp.status_code == 200)
    assert_contains("status ok", resp.text, '"status":"ok"')
    print()


def test_list_models(client):
    print("=== GET /v1/models ===")
    resp = client.get(f"{BASE}/v1/models")
    body = resp.text
    assert_contains("object is list", body, '"object":"list"')
    assert_contains("has data", body, '"data":')
    print()


def test_non_streaming_completion(client, model):
    print("=== POST /v1/chat/completions (non-streaming) ===")
    resp = client.post(
        f"{BASE}/v1/chat/completions",
        json=chat_request(model, stream=False),
    )
    body = resp.text
    assert_contains("object is chat.completion", body, '"object":"chat.completion"')
    assert_contains("has choices", body, '"choices":')
    assert_contains("has role assistant", body, '"role":"assistant"')
    assert_contains("has finish_reason", body, '"finish_reason":')

    data = resp.json()
    content = data["choices"][0]["message"]["content"]
    assert_true("content is non-empty", len(content) > 0)
    print()


def test_streaming_completion(client, model):
    print("=== POST /v1/chat/completions (streaming) ===")
    with client.stream(
        "POST",
        f"{BASE}/v1/chat/completions",
        json=chat_request(model, stream=True),
    ) as resp:
        lines = []
        for line in resp.iter_lines():
            if line:
                lines.append(line)

    full = "\n".join(lines)
    assert_contains("has SSE data lines", full, "data: ")
    assert_contains("object is chunk", full, '"object":"chat.completion.chunk"')
    assert_contains("has assistant role", full, '"role":"assistant"')

    # Last data line should have finish_reason
    data_lines = [l for l in lines if l.startswith("data: {")]
    assert_true("has at least 2 data lines", len(data_lines) >= 2)

    last_chunk = json.loads(data_lines[-1].removeprefix("data: "))
    finish_reason = last_chunk["choices"][0].get("finish_reason")
    assert_true(
        f"last chunk has finish_reason ({finish_reason})",
        finish_reason in ("stop", "length"),
    )

    # First chunk should have assistant role
    first_chunk = json.loads(data_lines[0].removeprefix("data: "))
    role = first_chunk["choices"][0]["delta"].get("role")
    assert_true("first chunk role is assistant", role == "assistant")
    print()


def test_max_tokens_1(client, model):
    print("=== Streaming with max_tokens=1 ===")
    with client.stream(
        "POST",
        f"{BASE}/v1/chat/completions",
        json=chat_request(model, stream=True, max_tokens=1, content="Tell me a long story."),
    ) as resp:
        lines = [line for line in resp.iter_lines() if line]

    data_lines = [l for l in lines if l.startswith("data: {")]
    assert_true("got data lines", len(data_lines) >= 1)

    last_chunk = json.loads(data_lines[-1].removeprefix("data: "))
    finish_reason = last_chunk["choices"][0].get("finish_reason")
    assert_true(
        f"has finish_reason ({finish_reason})",
        finish_reason in ("stop", "length"),
    )
    print()


def main():
    global PASS, FAIL

    with httpx.Client(timeout=60) as client:
        model = pick_model(client)
        print(f"Using model: {model}\n")

        test_health(client)
        test_list_models(client)
        test_non_streaming_completion(client, model)
        test_streaming_completion(client, model)
        test_max_tokens_1(client, model)

    print("=" * 30)
    print(f"Results: {PASS} passed, {FAIL} failed")
    if FAIL > 0:
        print("\033[31mSOME TESTS FAILED\033[0m")
        sys.exit(1)
    else:
        print("\033[32mALL TESTS PASSED\033[0m")


if __name__ == "__main__":
    main()