Files
web-page-backend/stock-lab/app/test_screener_engine.py
2026-05-12 09:29:10 +09:00

56 lines
2.0 KiB
Python

import datetime as dt
import pandas as pd
import pytest
from app.screener.engine import ScreenContext, Screener, combine
from app.screener.nodes.hygiene import HygieneGate
from app.screener.nodes.foreign_buy import ForeignBuy
from app.screener.nodes.momentum import Momentum20
from app.screener._test_fixtures import make_master, make_prices, make_flow, make_kospi
def _ctx(master, prices, flow):
return ScreenContext(master=master, prices=prices, flow=flow,
kospi=make_kospi(days=260),
asof=dt.date(2026, 5, 12))
def test_combine_weighted_average():
scores = {
"foreign_buy": pd.Series({"A": 80, "B": 20}),
"momentum": pd.Series({"A": 60, "B": 40}),
}
weights = {"foreign_buy": 2.0, "momentum": 1.0}
out = combine(scores, weights)
# A: (80*2 + 60*1)/3 = 73.33
assert abs(out["A"] - 73.333) < 0.1
assert abs(out["B"] - 26.666) < 0.1
def test_combine_all_zero_weight_raises():
scores = {"foreign_buy": pd.Series({"A": 80})}
with pytest.raises(ValueError, match="no active"):
combine(scores, {"foreign_buy": 0})
def test_screener_run_end_to_end():
asof = dt.date(2026, 5, 12)
master = make_master(["GOOD", "SMALL"],
market_caps={"GOOD": 200_000_000_000, "SMALL": 1_000_000_000})
prices = make_prices(["GOOD", "SMALL"], days=260, asof=asof, trend_pct=0.1)
flow = make_flow(["GOOD", "SMALL"], days=260, asof=asof,
foreign_per_day={"GOOD": 100_000_000, "SMALL": 0})
ctx = _ctx(master, prices, flow)
screener = Screener(
gate=HygieneGate(),
score_nodes=[ForeignBuy(), Momentum20()],
weights={"foreign_buy": 1.0, "momentum": 1.0},
node_params={"foreign_buy": {"window_days": 5}, "momentum": {"window_days": 20}},
gate_params={**HygieneGate.default_params, "min_listed_days": 0},
top_n=10,
)
result = screener.run(ctx)
assert result.survivors_count == 1 # SMALL은 게이트 탈락
assert result.ranked.index[0] == "GOOD"