mirror of
https://ghfast.top/https://github.com/aeroxw/tick-stock-panel.git
synced 2026-09-12 15:34:16 +08:00
103 lines
2.8 KiB
Python
103 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import pytest
|
|
|
|
from app.api.strategy import BuildRequest, build_strategy_stream
|
|
from app.strategy.ai_generator import AIStrategyGenerator
|
|
|
|
STREAM_CODE = '''"""测试策略"""
|
|
import polars as pl
|
|
|
|
META = {
|
|
"id": "wrong",
|
|
"name": "旧名",
|
|
"description": "旧描述",
|
|
"tags": ["测试"],
|
|
"params": [],
|
|
"scoring": {},
|
|
}
|
|
|
|
ENTRY_SIGNALS = []
|
|
EXIT_SIGNALS = []
|
|
STOP_LOSS = -0.05
|
|
MAX_HOLD_DAYS = 20
|
|
|
|
RULES = """
|
|
1. 测试规则一
|
|
2. 测试规则二
|
|
3. 测试规则三
|
|
"""
|
|
|
|
def filter(df: pl.DataFrame, params: dict) -> pl.Expr:
|
|
return pl.lit(True)
|
|
'''
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_build_strategy_stream_yields_delta_and_normalized_result(monkeypatch):
|
|
async def fake_stream(self, prompt):
|
|
yield STREAM_CODE[:40]
|
|
yield STREAM_CODE[40:]
|
|
|
|
monkeypatch.setattr(AIStrategyGenerator, "stream", fake_stream)
|
|
req = BuildRequest(
|
|
step=1,
|
|
name="新策略",
|
|
description="新描述",
|
|
direction="long",
|
|
rules="1. 规则一\n2. 规则二\n3. 规则三",
|
|
strategy_id="ai_streamed",
|
|
)
|
|
|
|
response = await build_strategy_stream(req, None)
|
|
body = b""
|
|
async for chunk in response.body_iterator:
|
|
body += chunk.encode("utf-8") if isinstance(chunk, str) else chunk
|
|
|
|
events = [json.loads(line) for line in body.decode("utf-8").splitlines()]
|
|
|
|
assert [event["type"] for event in events] == ["meta", "delta", "delta", "result"]
|
|
result = events[-1]
|
|
assert result["valid"] is True
|
|
assert result["meta"]["id"] == "ai_streamed"
|
|
assert result["meta"]["name"] == "新策略"
|
|
assert '"id": "ai_streamed"' in result["code"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_build_strategy_stream_repairs_missing_meta_once(monkeypatch):
|
|
calls = 0
|
|
|
|
async def fake_stream(self, prompt):
|
|
yield "import polars as pl\n\ndef filter(df, params):\n return pl.lit(True)\n"
|
|
|
|
async def fake_repair(self, code, error):
|
|
nonlocal calls
|
|
calls += 1
|
|
return self.validate_code(STREAM_CODE)
|
|
|
|
monkeypatch.setattr(AIStrategyGenerator, "stream", fake_stream)
|
|
monkeypatch.setattr(AIStrategyGenerator, "repair_code", fake_repair)
|
|
req = BuildRequest(
|
|
step=1,
|
|
name="修复后策略",
|
|
description="修复后描述",
|
|
direction="long",
|
|
rules="1. 测试规则",
|
|
strategy_id="ai_repaired",
|
|
)
|
|
|
|
response = await build_strategy_stream(req, None)
|
|
body = b""
|
|
async for chunk in response.body_iterator:
|
|
body += chunk.encode("utf-8") if isinstance(chunk, str) else chunk
|
|
|
|
result = json.loads(body.decode("utf-8").splitlines()[-1])
|
|
assert calls == 1
|
|
assert result["type"] == "result"
|
|
assert result["valid"] is True
|
|
assert result["meta"]["id"] == "ai_repaired"
|
|
assert result["meta"]["name"] == "修复后策略"
|