464 lines
18 KiB
Python
464 lines
18 KiB
Python
"""Tests for us2kr.collectors.prices (no network: yfinance.download / yfinance.Ticker are monkeypatched)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import date, datetime, timedelta
|
|
from typing import ClassVar
|
|
from zoneinfo import ZoneInfo
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
import yfinance as yf
|
|
|
|
from us2kr.collectors import prices
|
|
from us2kr.collectors.prices import PriceFetchError, fetch_quotes
|
|
|
|
SESSION = date(2026, 8, 24) # Monday
|
|
ET = ZoneInfo("America/New_York")
|
|
FIELDS = ["Open", "High", "Low", "Close", "Adj Close", "Volume"]
|
|
|
|
|
|
# ------------------------------------------------------------------------------------------------
|
|
# fixtures / helpers
|
|
# ------------------------------------------------------------------------------------------------
|
|
|
|
|
|
def _dates(n_before: int, n_after: int = 0) -> pd.DatetimeIndex:
|
|
"""Calendar-day index (weekends included, like the union index yf.download produces)."""
|
|
start = pd.Timestamp(SESSION) - pd.Timedelta(days=n_before)
|
|
return pd.date_range(start, pd.Timestamp(SESSION) + pd.Timedelta(days=n_after), freq="D")
|
|
|
|
|
|
def _ticker_block(idx: pd.DatetimeIndex, closes: dict[date, float], volume: float | None = 1_000.0):
|
|
"""Build a 6-column block for one ticker; rows not in ``closes`` are NaN (weekend / no trade)."""
|
|
block = pd.DataFrame(index=idx, columns=FIELDS, dtype="float64")
|
|
for d, c in closes.items():
|
|
ts = pd.Timestamp(d)
|
|
block.loc[ts, "Open"] = c - 1.0
|
|
block.loc[ts, "High"] = c + 2.0
|
|
block.loc[ts, "Low"] = c - 3.0
|
|
block.loc[ts, "Close"] = c
|
|
block.loc[ts, "Adj Close"] = c
|
|
block.loc[ts, "Volume"] = np.nan if volume is None else volume
|
|
return block
|
|
|
|
|
|
def _weekday_closes(last_close: float, n: int, step: float = 1.0, end: date = SESSION) -> dict[date, float]:
|
|
"""n weekday closes ending on ``end`` (inclusive), increasing by ``step`` per day: ... , last-step, last."""
|
|
out: dict[date, float] = {}
|
|
d = end
|
|
k = 0
|
|
while len(out) < n:
|
|
if d.weekday() < 5:
|
|
out[d] = last_close - step * k
|
|
k += 1
|
|
d -= timedelta(days=1)
|
|
return out
|
|
|
|
|
|
def _multi(blocks: dict[str, pd.DataFrame]) -> pd.DataFrame:
|
|
"""Assemble the (Ticker, Price) MultiIndex frame exactly like yf.download(group_by='ticker')."""
|
|
df = pd.concat(blocks.values(), axis=1, keys=blocks.keys(), names=["Ticker", "Price"])
|
|
df.index.name = "Date"
|
|
return df
|
|
|
|
|
|
class FakeTicker:
|
|
"""Stand-in for yfinance.Ticker; ``history`` returns pre-registered frames or raises."""
|
|
|
|
frames: ClassVar[dict[str, pd.DataFrame]] = {}
|
|
errors: ClassVar[set[str]] = set()
|
|
calls: ClassVar[list[tuple[str, dict]]] = []
|
|
|
|
def __init__(self, ticker: str, session=None):
|
|
self.ticker = ticker
|
|
|
|
def history(self, **kwargs):
|
|
FakeTicker.calls.append((self.ticker, kwargs))
|
|
if self.ticker in FakeTicker.errors:
|
|
raise RuntimeError(f"boom {self.ticker}")
|
|
return FakeTicker.frames.get(self.ticker, pd.DataFrame())
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _reset_fake_ticker(monkeypatch):
|
|
FakeTicker.frames = {}
|
|
FakeTicker.errors = set()
|
|
FakeTicker.calls = []
|
|
monkeypatch.setattr(yf, "Ticker", FakeTicker)
|
|
yield
|
|
|
|
|
|
@pytest.fixture
|
|
def batch_frame() -> pd.DataFrame:
|
|
idx = _dates(n_before=20, n_after=1) # includes the day AFTER the session (24h instruments)
|
|
gspc = _weekday_closes(6500.0, 12, step=10.0) # 6500 on SESSION, 6490 on Friday, ...
|
|
tnx = _weekday_closes(4.25, 12, step=-0.05) # 4.25 on SESSION, 4.30 on Friday (yield fell 5bp)
|
|
aapl = _weekday_closes(200.0, 12, step=2.0)
|
|
krw = _weekday_closes(1380.0, 12, step=5.0)
|
|
krw[SESSION + timedelta(days=1)] = 1399.0 # bar AFTER the session must be ignored
|
|
blocks = {
|
|
"^GSPC": _ticker_block(idx, gspc, volume=4.0e9),
|
|
"^TNX": _ticker_block(idx, tnx, volume=0.0),
|
|
"AAPL": _ticker_block(idx, aapl, volume=50e6),
|
|
"KRW=X": _ticker_block(idx, krw, volume=0.0),
|
|
"MISSING": _ticker_block(idx, {}), # all-NaN column, as yfinance emits for a bad symbol
|
|
}
|
|
# make AAPL's session volume differ from the prior average (prior 10 bars = 50e6, session = 100e6)
|
|
blocks["AAPL"].loc[pd.Timestamp(SESSION), "Volume"] = 100e6
|
|
return _multi(blocks)
|
|
|
|
|
|
@pytest.fixture
|
|
def patch_download(monkeypatch):
|
|
def _install(frame_or_exc, *sequence):
|
|
calls: list[dict] = []
|
|
items = [frame_or_exc, *sequence]
|
|
|
|
def fake_download(tickers, **kwargs):
|
|
calls.append({"tickers": list(tickers), **kwargs})
|
|
item = items[min(len(calls) - 1, len(items) - 1)]
|
|
if isinstance(item, Exception):
|
|
raise item
|
|
return item
|
|
|
|
monkeypatch.setattr(yf, "download", fake_download)
|
|
return calls
|
|
|
|
return _install
|
|
|
|
|
|
def _by_symbol(quotes):
|
|
return {q.symbol: q for q in quotes}
|
|
|
|
|
|
# ------------------------------------------------------------------------------------------------
|
|
# tests: per-group computations
|
|
# ------------------------------------------------------------------------------------------------
|
|
|
|
|
|
def test_batch_download_called_with_date_window(batch_frame, patch_download):
|
|
calls = patch_download(batch_frame)
|
|
fetch_quotes(SESSION, ["^GSPC", "AAPL"])
|
|
assert len(calls) == 1
|
|
call = calls[0]
|
|
assert call["tickers"] == ["^GSPC", "AAPL"]
|
|
assert call["group_by"] == "ticker"
|
|
assert call["auto_adjust"] is False
|
|
assert call["interval"] == "1d"
|
|
assert call["progress"] is False
|
|
assert call["start"] == (SESSION - timedelta(days=prices.HISTORY_DAYS)).isoformat()
|
|
assert call["end"] == (SESSION + timedelta(days=1)).isoformat() # end is exclusive
|
|
|
|
|
|
def test_index_quote_fields(batch_frame, patch_download):
|
|
patch_download(batch_frame)
|
|
q = _by_symbol(fetch_quotes(SESSION, ["^GSPC"]))["^GSPC"]
|
|
assert q.group == "indices"
|
|
assert q.name == "S&P 500"
|
|
assert q.close == pytest.approx(6500.0)
|
|
assert q.prev_close == pytest.approx(6490.0)
|
|
assert q.change_abs == pytest.approx(10.0)
|
|
assert q.change_pct == pytest.approx((6500.0 / 6490.0 - 1) * 100)
|
|
assert q.day_high == pytest.approx(6502.0)
|
|
assert q.day_low == pytest.approx(6497.0)
|
|
assert q.volume == pytest.approx(4.0e9)
|
|
assert q.avg_volume_10d == pytest.approx(4.0e9)
|
|
assert q.as_of == datetime(2026, 8, 24, tzinfo=ET)
|
|
assert q.source == "yfinance"
|
|
assert q.after_hours_price is None # indices never get after-hours lookups
|
|
|
|
|
|
def test_rates_change_abs_in_yield_points(batch_frame, patch_download):
|
|
patch_download(batch_frame)
|
|
q = _by_symbol(fetch_quotes(SESSION, ["^TNX"]))["^TNX"]
|
|
assert q.group == "rates"
|
|
assert q.close == pytest.approx(4.25)
|
|
assert q.prev_close == pytest.approx(4.30)
|
|
assert q.change_abs == pytest.approx(-0.05)
|
|
assert q.change_pct == pytest.approx((4.25 / 4.30 - 1) * 100)
|
|
assert q.volume is None # zero volume is reported as None
|
|
assert q.avg_volume_10d is None
|
|
|
|
|
|
def test_fx_close_selected_by_date_not_last_row(batch_frame, patch_download):
|
|
"""KRW=X has a bar dated after the session; close must be the session-date bar, prev the one before."""
|
|
patch_download(batch_frame)
|
|
q = _by_symbol(fetch_quotes(SESSION, ["KRW=X"]))["KRW=X"]
|
|
assert q.group == "fx"
|
|
assert q.close == pytest.approx(1380.0)
|
|
assert q.prev_close == pytest.approx(1375.0)
|
|
assert q.change_abs == pytest.approx(5.0)
|
|
|
|
|
|
def test_equity_volume_ratio_inputs(batch_frame, patch_download):
|
|
patch_download(batch_frame)
|
|
q = _by_symbol(fetch_quotes(SESSION, ["AAPL"]))["AAPL"]
|
|
assert q.group == "megacap"
|
|
assert q.name == "애플"
|
|
assert q.volume == pytest.approx(100e6)
|
|
assert q.avg_volume_10d == pytest.approx(50e6) # prior 10 sessions only, session day excluded
|
|
|
|
|
|
def test_missing_ticker_skipped_and_order_preserved(batch_frame, patch_download, caplog):
|
|
patch_download(batch_frame)
|
|
with caplog.at_level("WARNING", logger="us2kr.collectors.prices"):
|
|
quotes = fetch_quotes(SESSION, ["^GSPC", "MISSING", "AAPL", "NOTINFRAME"])
|
|
assert [q.symbol for q in quotes] == ["^GSPC", "AAPL"]
|
|
assert "MISSING" in caplog.text
|
|
assert "NOTINFRAME" in caplog.text
|
|
|
|
|
|
def test_unknown_ticker_gets_other_group(patch_download):
|
|
idx = _dates(10)
|
|
frame = _multi({"ZZZ": _ticker_block(idx, _weekday_closes(10.0, 5))})
|
|
patch_download(frame)
|
|
q = fetch_quotes(SESSION, ["ZZZ"])[0]
|
|
assert q.group == "other"
|
|
assert q.name == "ZZZ"
|
|
|
|
|
|
def test_default_tickers_is_symbol_universe(batch_frame, patch_download):
|
|
calls = patch_download(batch_frame)
|
|
fetch_quotes(SESSION)
|
|
from us2kr import symbols
|
|
|
|
assert calls[0]["tickers"] == symbols.all_tickers()
|
|
|
|
|
|
def test_empty_ticker_list_returns_empty_without_download(patch_download):
|
|
calls = patch_download(pd.DataFrame())
|
|
assert fetch_quotes(SESSION, []) == []
|
|
assert calls == []
|
|
|
|
|
|
# ------------------------------------------------------------------------------------------------
|
|
# tests: session_date edge cases
|
|
# ------------------------------------------------------------------------------------------------
|
|
|
|
|
|
def test_session_row_missing_uses_last_prior_bar(patch_download, caplog):
|
|
"""If the session-date bar is absent (not yet published), fall back to the last bar before it."""
|
|
idx = _dates(10)
|
|
friday = SESSION - timedelta(days=3)
|
|
closes = _weekday_closes(100.0, 5, step=1.0, end=friday) # 100 on Fri, 99 on Thu
|
|
frame = _multi({"XLK": _ticker_block(idx, closes)})
|
|
patch_download(frame)
|
|
with caplog.at_level("INFO", logger="us2kr.collectors.prices"):
|
|
q = fetch_quotes(SESSION, ["XLK"])[0]
|
|
assert q.close == pytest.approx(100.0)
|
|
assert q.prev_close == pytest.approx(99.0)
|
|
assert q.as_of.date() == friday
|
|
assert "no 2026-08-24 bar" in caplog.text
|
|
|
|
|
|
def test_only_one_bar_gives_no_change(patch_download):
|
|
idx = _dates(3)
|
|
frame = _multi({"XLF": _ticker_block(idx, {SESSION: 50.0})})
|
|
patch_download(frame)
|
|
q = fetch_quotes(SESSION, ["XLF"])[0]
|
|
assert q.close == pytest.approx(50.0)
|
|
assert q.prev_close is None
|
|
assert q.change_pct is None
|
|
assert q.change_abs is None
|
|
assert q.avg_volume_10d is None
|
|
|
|
|
|
def test_nan_close_rows_are_ignored_for_prev_close(patch_download):
|
|
idx = _dates(10)
|
|
closes = _weekday_closes(100.0, 5, step=1.0)
|
|
block = _ticker_block(idx, closes)
|
|
friday = SESSION - timedelta(days=3)
|
|
block.loc[pd.Timestamp(friday), "Close"] = np.nan # broken row must be skipped, not used as prev
|
|
frame = _multi({"XLE": block})
|
|
patch_download(frame)
|
|
q = fetch_quotes(SESSION, ["XLE"])[0]
|
|
assert q.close == pytest.approx(100.0)
|
|
assert q.prev_close == pytest.approx(98.0) # Thursday
|
|
|
|
|
|
def test_nan_volume_becomes_none(patch_download):
|
|
idx = _dates(10)
|
|
frame = _multi({"XLV": _ticker_block(idx, _weekday_closes(100.0, 5), volume=None)})
|
|
patch_download(frame)
|
|
q = fetch_quotes(SESSION, ["XLV"])[0]
|
|
assert q.volume is None
|
|
assert q.avg_volume_10d is None
|
|
|
|
|
|
def test_all_bars_after_session_date_skips_ticker(patch_download):
|
|
idx = _dates(0, n_after=2)
|
|
frame = _multi({"BTC-USD": _ticker_block(idx, {SESSION + timedelta(days=1): 1.0})})
|
|
patch_download(frame)
|
|
assert fetch_quotes(SESSION, ["BTC-USD"]) == []
|
|
|
|
|
|
# ------------------------------------------------------------------------------------------------
|
|
# tests: frame layouts
|
|
# ------------------------------------------------------------------------------------------------
|
|
|
|
|
|
def test_single_level_frame_for_single_ticker(patch_download):
|
|
idx = _dates(10)
|
|
frame = _ticker_block(idx, _weekday_closes(30.0, 5, step=0.5))
|
|
patch_download(frame)
|
|
q = fetch_quotes(SESSION, ["^SOX"])[0]
|
|
assert q.close == pytest.approx(30.0)
|
|
assert q.prev_close == pytest.approx(29.5)
|
|
|
|
|
|
def test_group_by_column_layout_is_handled(batch_frame, patch_download):
|
|
swapped = batch_frame.swaplevel(0, 1, axis=1).sort_index(axis=1)
|
|
patch_download(swapped)
|
|
q = _by_symbol(fetch_quotes(SESSION, ["^GSPC"]))["^GSPC"]
|
|
assert q.close == pytest.approx(6500.0)
|
|
|
|
|
|
def test_tz_aware_index_is_normalised(patch_download):
|
|
idx = pd.date_range(pd.Timestamp(SESSION) - pd.Timedelta(days=6), pd.Timestamp(SESSION), freq="D")
|
|
block = _ticker_block(idx, _weekday_closes(10.0, 4))
|
|
block.index = block.index.tz_localize(ET)
|
|
patch_download(_multi({"SMH": block}))
|
|
q = fetch_quotes(SESSION, ["SMH"])[0]
|
|
assert q.close == pytest.approx(10.0)
|
|
assert q.prev_close == pytest.approx(9.0)
|
|
|
|
|
|
@pytest.mark.parametrize("tz", ["Europe/London", "UTC", "America/New_York", "America/Chicago"])
|
|
def test_tz_aware_daily_index_keeps_exchange_local_date(tz):
|
|
"""London/UTC-stamped daily bars (KRW=X, BTC-USD via Ticker.history) must not shift a day earlier."""
|
|
idx = pd.DatetimeIndex(pd.to_datetime(["2026-08-20", "2026-08-21", "2026-08-24", "2026-08-25"]))
|
|
frame = pd.DataFrame(
|
|
{"Close": [1380.0, 1385.0, 1390.0, 1400.0], "Volume": [0.0] * 4}, index=idx.tz_localize(tz)
|
|
)
|
|
cleaned = prices._clean_single(frame)
|
|
assert cleaned is not None and cleaned.index.tz is None
|
|
dates = [d.date().isoformat() for d in cleaned.index]
|
|
assert dates == ["2026-08-20", "2026-08-21", "2026-08-24", "2026-08-25"]
|
|
q = prices._quote_from_history("KRW=X", cleaned, SESSION)
|
|
assert q.close == pytest.approx(1390.0)
|
|
assert q.prev_close == pytest.approx(1385.0)
|
|
assert q.as_of == datetime(2026, 8, 24, tzinfo=ET)
|
|
|
|
|
|
# ------------------------------------------------------------------------------------------------
|
|
# tests: failure policy
|
|
# ------------------------------------------------------------------------------------------------
|
|
|
|
|
|
def test_batch_retried_once_then_succeeds(batch_frame, patch_download):
|
|
calls = patch_download(RuntimeError("yahoo down"), batch_frame)
|
|
quotes = fetch_quotes(SESSION, ["^GSPC"])
|
|
assert len(calls) == 2
|
|
assert quotes[0].close == pytest.approx(6500.0)
|
|
|
|
|
|
def test_batch_failure_falls_back_to_indices_only(patch_download):
|
|
patch_download(RuntimeError("down"), RuntimeError("still down"))
|
|
idx = _dates(10)
|
|
FakeTicker.frames["^GSPC"] = _ticker_block(idx, _weekday_closes(6000.0, 5, step=10.0))
|
|
FakeTicker.frames["AAPL"] = _ticker_block(idx, _weekday_closes(200.0, 5))
|
|
quotes = fetch_quotes(SESSION, ["^GSPC", "AAPL", "^VIX"])
|
|
# only indices are attempted in fallback; AAPL is not requested, ^VIX yields nothing
|
|
assert [q.symbol for q in quotes] == ["^GSPC"]
|
|
assert quotes[0].close == pytest.approx(6000.0)
|
|
history_calls = [c for c in FakeTicker.calls if c[1].get("interval") == "1d"]
|
|
assert sorted(t for t, _ in history_calls) == ["^GSPC", "^VIX"]
|
|
|
|
|
|
def test_empty_batch_frames_count_as_failure(patch_download):
|
|
patch_download(pd.DataFrame(), None)
|
|
FakeTicker.frames["^DJI"] = _ticker_block(_dates(10), _weekday_closes(40000.0, 5))
|
|
quotes = fetch_quotes(SESSION, ["^DJI"])
|
|
assert quotes[0].close == pytest.approx(40000.0)
|
|
|
|
|
|
def test_total_failure_raises_price_fetch_error(patch_download):
|
|
patch_download(RuntimeError("down"))
|
|
FakeTicker.errors = {"^GSPC"}
|
|
with pytest.raises(PriceFetchError) as exc_info:
|
|
fetch_quotes(SESSION, ["^GSPC", "AAPL"])
|
|
assert "down" in str(exc_info.value)
|
|
assert "^GSPC" in str(exc_info.value)
|
|
|
|
|
|
def test_run_safely_records_gap_on_total_failure(patch_download):
|
|
from us2kr.collectors.base import run_safely
|
|
|
|
patch_download(RuntimeError("down"))
|
|
gaps: list[str] = []
|
|
out = run_safely("prices", lambda: fetch_quotes(SESSION, ["AAPL"]), [], gaps)
|
|
assert out == []
|
|
assert gaps and gaps[0].startswith("prices: PriceFetchError")
|
|
|
|
|
|
# ------------------------------------------------------------------------------------------------
|
|
# tests: after-hours (best effort)
|
|
# ------------------------------------------------------------------------------------------------
|
|
|
|
|
|
def _minute_frame(closes_by_time: dict[str, float], day: date = SESSION) -> pd.DataFrame:
|
|
stamps = [pd.Timestamp(f"{day.isoformat()} {t}", tz=ET) for t in closes_by_time]
|
|
return pd.DataFrame(
|
|
{"Open": list(closes_by_time.values()), "Close": list(closes_by_time.values()), "Volume": 0},
|
|
index=pd.DatetimeIndex(stamps, name="Datetime"),
|
|
)
|
|
|
|
|
|
def test_after_hours_price_from_post_close_minute_bars(batch_frame, patch_download):
|
|
patch_download(batch_frame)
|
|
FakeTicker.frames["AAPL"] = _minute_frame(
|
|
{"15:58": 199.0, "15:59": 200.0, "16:00": 201.0, "19:59": 204.0}
|
|
)
|
|
q = _by_symbol(fetch_quotes(SESSION, ["AAPL", "^GSPC", "KRW=X"]))
|
|
assert q["AAPL"].after_hours_price == pytest.approx(204.0)
|
|
assert q["AAPL"].after_hours_change_pct == pytest.approx(2.0)
|
|
assert q["^GSPC"].after_hours_price is None
|
|
assert q["KRW=X"].after_hours_price is None
|
|
# only the equity was queried, with prepost 1m bars for the session day
|
|
assert [t for t, _ in FakeTicker.calls] == ["AAPL"]
|
|
kwargs = FakeTicker.calls[0][1]
|
|
assert kwargs["interval"] == "1m" and kwargs["prepost"] is True
|
|
assert kwargs["start"] == SESSION.isoformat()
|
|
|
|
|
|
def test_after_hours_skipped_when_only_regular_bars(batch_frame, patch_download):
|
|
patch_download(batch_frame)
|
|
FakeTicker.frames["AAPL"] = _minute_frame({"15:58": 199.0, "15:59": 200.0})
|
|
q = _by_symbol(fetch_quotes(SESSION, ["AAPL"]))["AAPL"]
|
|
assert q.after_hours_price is None
|
|
assert q.after_hours_change_pct is None
|
|
|
|
|
|
def test_after_hours_ignores_bars_from_other_days(batch_frame, patch_download):
|
|
patch_download(batch_frame)
|
|
FakeTicker.frames["AAPL"] = _minute_frame({"17:00": 150.0}, day=SESSION - timedelta(days=3))
|
|
q = _by_symbol(fetch_quotes(SESSION, ["AAPL"]))["AAPL"]
|
|
assert q.after_hours_price is None
|
|
|
|
|
|
def test_after_hours_failure_is_isolated(batch_frame, patch_download):
|
|
patch_download(batch_frame)
|
|
FakeTicker.errors = {"AAPL"}
|
|
q = _by_symbol(fetch_quotes(SESSION, ["AAPL", "^GSPC"]))
|
|
assert q["AAPL"].close == pytest.approx(200.0)
|
|
assert q["AAPL"].after_hours_price is None
|
|
assert q["^GSPC"].close == pytest.approx(6500.0)
|
|
|
|
|
|
def test_after_hours_respects_time_budget(batch_frame, patch_download, monkeypatch):
|
|
patch_download(batch_frame)
|
|
FakeTicker.frames["AAPL"] = _minute_frame({"16:30": 210.0})
|
|
monkeypatch.setattr(prices, "AFTER_HOURS_BUDGET_SECONDS", 0.0)
|
|
q = _by_symbol(fetch_quotes(SESSION, ["AAPL"]))["AAPL"]
|
|
assert q.after_hours_price is None
|
|
assert FakeTicker.calls == []
|
|
|
|
|
|
def test_is_equity_like_classification():
|
|
assert prices._is_equity_like("AAPL")
|
|
assert prices._is_equity_like("EWY")
|
|
for t in ("^GSPC", "ES=F", "KRW=X", "DX-Y.NYB", "BTC-USD", "^TNX"):
|
|
assert not prices._is_equity_like(t), t
|