221 lines
9.6 KiB
Python
221 lines
9.6 KiB
Python
"""
|
|
Regression tests for scripts/serve_openai_api.py
|
|
These tests verify OpenAI API compatibility without loading a real model.
|
|
"""
|
|
import json
|
|
import re
|
|
import sys
|
|
import os
|
|
import types
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Stub heavy dependencies so importing serve_openai_api doesn't need GPU/model
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# Provide a minimal torch stub (only attributes actually used at import time)
|
|
_torch_stub = types.ModuleType("torch")
|
|
_torch_stub.cuda = types.SimpleNamespace(is_available=lambda: False)
|
|
_torch_stub.no_grad = lambda: types.SimpleNamespace(__enter__=lambda s: None, __exit__=lambda s, *a: None)
|
|
_torch_stub.inference_mode = lambda: lambda fn: fn
|
|
_torch_stub.load = lambda *a, **kw: {}
|
|
_torch_stub.Tensor = type("Tensor", (), {})
|
|
_torch_stub.float32 = "float32"
|
|
_torch_stub.LongTensor = lambda *a, **kw: None
|
|
|
|
# optim sub-module
|
|
_optim_stub = types.ModuleType("torch.optim")
|
|
sys.modules["torch.optim"] = _optim_stub
|
|
_torch_stub.optim = _optim_stub
|
|
|
|
# nn sub-module with Module and common layers
|
|
_nn_stub = types.ModuleType("torch.nn")
|
|
class _ModuleStub:
|
|
def __init_subclass__(cls, **kw): pass
|
|
def __init__(self, *a, **kw): pass
|
|
_nn_stub.Module = _ModuleStub
|
|
_nn_stub.Linear = type("Linear", (_ModuleStub,), {})
|
|
_nn_stub.Embedding = type("Embedding", (_ModuleStub,), {})
|
|
_nn_stub.Dropout = type("Dropout", (_ModuleStub,), {})
|
|
class _ParameterStub:
|
|
def __init__(self, *a, **kw): pass
|
|
_nn_stub.Parameter = _ParameterStub
|
|
_torch_stub.nn = _nn_stub
|
|
|
|
_nn_func = types.ModuleType("torch.nn.functional")
|
|
_nn_func.scaled_dot_product_attention = lambda *a, **kw: None
|
|
_nn_func.cross_entropy = lambda *a, **kw: None
|
|
_nn_func.softmax = lambda *a, **kw: None
|
|
_nn_func.silu = lambda *a, **kw: None
|
|
|
|
sys.modules["torch"] = _torch_stub
|
|
sys.modules["torch.nn"] = _nn_stub
|
|
sys.modules["torch.nn.functional"] = _nn_func
|
|
sys.modules.setdefault("torch.distributed", types.ModuleType("torch.distributed"))
|
|
sys.modules.setdefault("torch.nn.parallel", types.ModuleType("torch.nn.parallel"))
|
|
sys.modules.setdefault("torch.utils", types.ModuleType("torch.utils"))
|
|
sys.modules.setdefault("torch.utils.data", types.ModuleType("torch.utils.data"))
|
|
|
|
for mod_name in [
|
|
"transformers", "transformers.activations", "transformers.modeling_outputs",
|
|
"librosa", "soundfile", "numpy", "uvicorn",
|
|
]:
|
|
sys.modules.setdefault(mod_name, types.ModuleType(mod_name))
|
|
|
|
# Stub transformers classes that are used at module scope
|
|
_transformers = sys.modules["transformers"]
|
|
for cls_name in [
|
|
"PreTrainedModel", "GenerationMixin", "PretrainedConfig",
|
|
"AutoTokenizer", "AutoModelForCausalLM", "AutoModel",
|
|
"AutoModelForSequenceClassification", "TextStreamer",
|
|
]:
|
|
if not hasattr(_transformers, cls_name):
|
|
setattr(_transformers, cls_name, type(cls_name, (), {}))
|
|
|
|
_mo = sys.modules["transformers.modeling_outputs"]
|
|
if not hasattr(_mo, "MoeCausalLMOutputWithPast"):
|
|
setattr(_mo, "MoeCausalLMOutputWithPast", type("MoeCausalLMOutputWithPast", (), {}))
|
|
|
|
_act = sys.modules["transformers.activations"]
|
|
if not hasattr(_act, "ACT2FN"):
|
|
setattr(_act, "ACT2FN", {})
|
|
|
|
# Now import the functions we can test without a live model
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "scripts"))
|
|
|
|
from scripts.serve_openai_api import parse_response
|
|
|
|
|
|
# ============================================================
|
|
# Tests for parse_response
|
|
# ============================================================
|
|
|
|
class TestParseResponse:
|
|
"""Verify that parse_response correctly separates reasoning, content and tool calls."""
|
|
|
|
def test_plain_text(self):
|
|
content, reasoning, tool_calls = parse_response("Hello, world!")
|
|
assert content == "Hello, world!"
|
|
assert reasoning is None
|
|
assert tool_calls is None
|
|
|
|
def test_think_tags_complete(self):
|
|
text = "<think>Step 1: consider options</think>The answer is 42."
|
|
content, reasoning, tool_calls = parse_response(text)
|
|
assert reasoning == "Step 1: consider options"
|
|
assert content == "The answer is 42."
|
|
assert tool_calls is None
|
|
|
|
def test_think_tag_no_opening(self):
|
|
"""When only </think> is present (streaming partial), split at the tag."""
|
|
text = "I need to think carefully</think>Final answer here"
|
|
content, reasoning, tool_calls = parse_response(text)
|
|
assert reasoning == "I need to think carefully"
|
|
assert content == "Final answer here"
|
|
|
|
def test_tool_call_parsed(self):
|
|
text = 'Sure, I can help. <tool_call>{"name":"get_weather","arguments":{"city":"Tokyo"}}</tool_call>'
|
|
content, reasoning, tool_calls = parse_response(text)
|
|
assert "Sure, I can help." in content
|
|
assert tool_calls is not None
|
|
assert len(tool_calls) == 1
|
|
assert tool_calls[0]["function"]["name"] == "get_weather"
|
|
args = json.loads(tool_calls[0]["function"]["arguments"])
|
|
assert args["city"] == "Tokyo"
|
|
|
|
def test_multiple_tool_calls(self):
|
|
text = (
|
|
'<tool_call>{"name":"a","arguments":{}}</tool_call>'
|
|
'<tool_call>{"name":"b","arguments":{"x":1}}</tool_call>'
|
|
)
|
|
content, reasoning, tool_calls = parse_response(text)
|
|
assert tool_calls is not None
|
|
assert len(tool_calls) == 2
|
|
assert tool_calls[0]["function"]["name"] == "a"
|
|
assert tool_calls[1]["function"]["name"] == "b"
|
|
|
|
def test_invalid_tool_call_json_skipped(self):
|
|
text = '<tool_call>not valid json</tool_call>The rest.'
|
|
content, reasoning, tool_calls = parse_response(text)
|
|
assert "The rest." in content
|
|
assert tool_calls is None # Invalid JSON is silently skipped
|
|
|
|
def test_empty_string(self):
|
|
content, reasoning, tool_calls = parse_response("")
|
|
assert content == ""
|
|
assert reasoning is None
|
|
assert tool_calls is None
|
|
|
|
|
|
# ============================================================
|
|
# Tests for SSE stream format (structural, no live model)
|
|
# ============================================================
|
|
|
|
class TestSSEStreamFormat:
|
|
"""Verify the SSE event stream wrapper produces correct framing."""
|
|
|
|
def test_done_marker_present(self):
|
|
"""The stream MUST end with 'data: [DONE]\\n\\n' per OpenAI spec."""
|
|
# Read the source and verify the [DONE] marker is emitted
|
|
script_path = os.path.join(os.path.dirname(__file__), "..", "scripts", "serve_openai_api.py")
|
|
with open(script_path) as f:
|
|
source = f.read()
|
|
assert 'data: [DONE]' in source, "Stream must emit 'data: [DONE]' terminator"
|
|
|
|
def test_stream_chunks_have_required_fields(self):
|
|
"""Each stream chunk must include id, object, created, model, choices."""
|
|
script_path = os.path.join(os.path.dirname(__file__), "..", "scripts", "serve_openai_api.py")
|
|
with open(script_path) as f:
|
|
source = f.read()
|
|
# Verify _make_chunk builds proper structure with all required fields
|
|
assert '"id": request_id' in source or '"id":' in source
|
|
assert '"object": "chat.completion.chunk"' in source
|
|
assert '"created":' in source
|
|
assert '"model":' in source
|
|
|
|
def test_cors_middleware_present(self):
|
|
"""CORS middleware must be configured for browser-based clients."""
|
|
script_path = os.path.join(os.path.dirname(__file__), "..", "scripts", "serve_openai_api.py")
|
|
with open(script_path) as f:
|
|
source = f.read()
|
|
assert "CORSMiddleware" in source, "CORS middleware required for browser clients"
|
|
|
|
def test_non_stream_uses_max_new_tokens(self):
|
|
"""Non-stream path must use max_new_tokens, not max_length, matching stream behavior."""
|
|
script_path = os.path.join(os.path.dirname(__file__), "..", "scripts", "serve_openai_api.py")
|
|
with open(script_path) as f:
|
|
source = f.read()
|
|
# The non-stream generate call should use max_new_tokens
|
|
# Find the non-stream generate block (after "with torch.no_grad():")
|
|
non_stream_section = source[source.index("with torch.no_grad()"):]
|
|
assert "max_new_tokens" in non_stream_section, \
|
|
"Non-stream path should use max_new_tokens for consistent behavior with stream path"
|
|
# max_length should NOT appear in the generate call
|
|
generate_block = non_stream_section[:non_stream_section.index("answer = tokenizer.decode")]
|
|
assert "max_length" not in generate_block, \
|
|
"Non-stream path should not use max_length (use max_new_tokens instead)"
|
|
|
|
def test_generation_thread_is_daemon(self):
|
|
"""Generation thread must be a daemon so it doesn't block process exit."""
|
|
script_path = os.path.join(os.path.dirname(__file__), "..", "scripts", "serve_openai_api.py")
|
|
with open(script_path) as f:
|
|
source = f.read()
|
|
assert "daemon=True" in source, "Generation thread should be daemon to avoid blocking exit"
|
|
|
|
def test_generation_thread_error_propagation(self):
|
|
"""Generation thread errors must be propagated to the consumer, not swallowed."""
|
|
script_path = os.path.join(os.path.dirname(__file__), "..", "scripts", "serve_openai_api.py")
|
|
with open(script_path) as f:
|
|
source = f.read()
|
|
# The _generate() function should have a try/except that puts errors on the queue
|
|
assert "_SENTINEL" in source, "Sentinel-based error propagation required"
|
|
|
|
|
|
# ============================================================
|
|
# Run with pytest or directly
|
|
# ============================================================
|
|
|
|
if __name__ == "__main__":
|
|
import pytest
|
|
sys.exit(pytest.main([__file__, "-v"]))
|