62 lines
2.1 KiB
Python
62 lines
2.1 KiB
Python
import datetime as dt
|
||
import sqlite3
|
||
|
||
import pandas as pd
|
||
import pytest
|
||
|
||
from app.screener.engine import ScreenContext
|
||
from app.screener.schema import ensure_screener_schema
|
||
from app.screener._test_fixtures import make_master, make_prices, make_flow
|
||
|
||
|
||
@pytest.fixture
|
||
def conn(tmp_path):
|
||
db_path = tmp_path / "ctx.db"
|
||
c = sqlite3.connect(db_path)
|
||
ensure_screener_schema(c)
|
||
yield c
|
||
c.close()
|
||
|
||
|
||
def _seed(conn, master_df, prices_df, flow_df):
|
||
now = dt.datetime.utcnow().isoformat()
|
||
for t, row in master_df.iterrows():
|
||
conn.execute("""INSERT INTO krx_master (ticker,name,market,market_cap,
|
||
is_managed,is_preferred,is_spac,listed_date,updated_at)
|
||
VALUES (?,?,?,?,?,?,?,?,?)""",
|
||
(t, row["name"], row["market"], row["market_cap"],
|
||
row["is_managed"], row["is_preferred"], row["is_spac"], None, now))
|
||
prices_df.to_sql("krx_daily_prices", conn, if_exists="append", index=False)
|
||
flow_df.to_sql("krx_flow", conn, if_exists="append", index=False)
|
||
conn.commit()
|
||
|
||
|
||
def test_load_returns_dataframes(conn):
|
||
asof = dt.date(2026, 5, 12)
|
||
_seed(conn,
|
||
make_master(["005930", "035420"]),
|
||
make_prices(["005930", "035420"], days=30, asof=asof),
|
||
make_flow(["005930", "035420"], days=30, asof=asof))
|
||
|
||
ctx = ScreenContext.load(conn, asof, lookback_days=30)
|
||
|
||
assert ctx.asof == asof
|
||
assert set(ctx.master.index) == {"005930", "035420"}
|
||
assert ctx.prices.shape[0] == 60 # 2 종목 × 30일
|
||
assert ctx.flow.shape[0] == 60
|
||
|
||
|
||
def test_restrict_filters_tickers(conn):
|
||
asof = dt.date(2026, 5, 12)
|
||
_seed(conn,
|
||
make_master(["005930", "035420", "091990"]),
|
||
make_prices(["005930", "035420", "091990"], days=30, asof=asof),
|
||
make_flow(["005930", "035420", "091990"], days=30, asof=asof))
|
||
|
||
ctx = ScreenContext.load(conn, asof, lookback_days=30)
|
||
scoped = ctx.restrict(pd.Index(["005930"]))
|
||
|
||
assert list(scoped.master.index) == ["005930"]
|
||
assert (scoped.prices["ticker"] == "005930").all()
|
||
assert (scoped.flow["ticker"] == "005930").all()
|