ViperEkura
2196c34c52
- 新增 protocol.py 协议层，Template Method 模式消除流/非流分支 45% 重复 - SSEBuilder 统一 SSE 构造，StopChecker 独立 stop_sequence 检测 - AnthropicHandler 追踪已产出文本，修复 stop 时重复 delta - server.py 路由从约 100 行缩减至 3 行 - 拆分为 core/（cache/executor/scheduler/task）和 api/（protocol/server） - 外部保持二级导入路径（from astrai.inference import Name） - 删除所有分隔线注释，代码按语义自然分组
54 lines
1.3 KiB
Python
54 lines
1.3 KiB
Python
"""Shared fixtures for inference tests."""
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
from astrai.inference import app
|
|
|
|
|
|
@pytest.fixture
|
|
def client():
|
|
"""Provide a test client for the FastAPI app."""
|
|
app.state.server_config = {
|
|
"device": "cpu",
|
|
"dtype": "bfloat16",
|
|
"param_path": None,
|
|
"max_batch_size": 1,
|
|
"_test": True,
|
|
}
|
|
app.state.engine = None
|
|
return TestClient(app)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_engine():
|
|
"""Create a mock InferenceEngine."""
|
|
|
|
async def _async_gen():
|
|
yield "chunk1"
|
|
yield "chunk2"
|
|
yield "[DONE]"
|
|
|
|
mock = MagicMock()
|
|
mock.generate.return_value = "mock response"
|
|
mock.generate_async.return_value = _async_gen()
|
|
mock.get_stats.return_value = {
|
|
"total_tasks": 0,
|
|
"total_tokens": 0,
|
|
"active_tasks": 0,
|
|
"waiting_queue": 0,
|
|
}
|
|
mock.tokenizer.encode.return_value = [1, 2, 3]
|
|
mock.tokenizer.decode.return_value = "mock response"
|
|
mock.tokenizer.apply_chat_template.return_value = "mock prompt"
|
|
return mock
|
|
|
|
|
|
@pytest.fixture
|
|
def loaded_model(client, mock_engine):
|
|
"""Simulate that the engine is loaded."""
|
|
app.state.engine = mock_engine
|
|
return mock_engine
|