diff --git a/collectors/universe_returns.py b/collectors/universe_returns.py index f55c6a4..ff34fa0 100644 --- a/collectors/universe_returns.py +++ b/collectors/universe_returns.py @@ -23,6 +23,7 @@ import boto3 import pandas as pd +from alpha_engine_lib.trading_calendar import add_trading_days as _add_trading_days logger = logging.getLogger(__name__) @@ -224,7 +225,7 @@ def _trading_days_to_process( if nyse_is_trading_day(d): trading_days_seen += 1 iso = d.isoformat() - fwd_5d = _add_business_days(d, 5) + fwd_5d = _add_trading_days(d, 5) if fwd_5d < today and iso not in existing: out.append(iso) d -= timedelta(days=1) @@ -328,9 +329,9 @@ def _build_rows_for_date( ) -> list[dict]: """Build universe_returns rows for a single eval_date.""" eval_dt = date.fromisoformat(eval_date) - fwd_5d = _add_business_days(eval_dt, 5) - fwd_10d = _add_business_days(eval_dt, 10) - fwd_30d = _add_business_days(eval_dt, 30) + fwd_5d = _add_trading_days(eval_dt, 5) + fwd_10d = _add_trading_days(eval_dt, 10) + fwd_30d = _add_trading_days(eval_dt, 30) # Check that forward dates are in the past (returns can be computed) today = date.today() @@ -350,7 +351,7 @@ def _build_rows_for_date( if not prices_t0: logger.warning("No prices for eval_date %s — may be a non-trading day", eval_date) # Try next business day - next_day = _add_business_days(eval_dt, 1) + next_day = _add_trading_days(eval_dt, 1) prices_t0 = polygon_client.get_grouped_daily(str(next_day)) if not prices_t0: return [] @@ -424,14 +425,3 @@ def _pct_return(price_start: float | None, price_end: float | None) -> float | N if price_start is None or price_end is None or price_start <= 0: return None return (price_end / price_start) - 1.0 - - -def _add_business_days(start: date, n: int) -> date: - """Add n business days to a date (skipping weekends).""" - current = start - added = 0 - while added < n: - current += timedelta(days=1) - if current.weekday() < 5: # Mon-Fri - added += 1 - return current diff --git a/tests/test_universe_returns_trading_day.py b/tests/test_universe_returns_trading_day.py new file mode 100644 index 0000000..d3f44c7 --- /dev/null +++ b/tests/test_universe_returns_trading_day.py @@ -0,0 +1,73 @@ +"""Trading-day eligibility filter in collectors/universe_returns.py. + +The arithmetic itself (add_trading_days) is locked in alpha-engine-lib's +test_trading_calendar.py. These tests cover the data-module behavior: +that _trading_days_to_process correctly delegates to NYSE-aware arithmetic +and never enqueues weekends, holidays, or eval_dates whose 5d forward +window has not yet closed. +""" +from __future__ import annotations + +from datetime import date + +from collectors.universe_returns import _trading_days_to_process + + +class TestTradingDaysToProcess: + def test_excludes_weekends_and_holidays(self): + # Today = Sat 2026-04-25; lookback covers Good Friday 4/3 and the + # 4/4-5 weekend. None should appear in the result. + out = _trading_days_to_process( + today=date(2026, 4, 25), + max_lookback=20, + existing=set(), + ) + forbidden = {"2026-04-03", "2026-04-04", "2026-04-05"} + assert forbidden.isdisjoint(set(out)) + + def test_eligibility_requires_fwd_5d_strictly_past(self): + # Today = 2026-04-25 (Sat). 4/20 (Mon) has fwd_5d = 4/27 (next Mon), + # which is NOT < today → ineligible. + out = _trading_days_to_process( + today=date(2026, 4, 25), + max_lookback=10, + existing=set(), + ) + assert "2026-04-20" not in out + + def test_existing_dates_filtered(self): + out = _trading_days_to_process( + today=date(2026, 4, 25), + max_lookback=20, + existing={"2026-04-17"}, + ) + assert "2026-04-17" not in out + + def test_4_17_eligible_today(self): + # On 2026-04-25, 2026-04-17 (Fri) IS eligible because + # fwd_5d = 4/24 < 4/25 and it's a trading day. + out = _trading_days_to_process( + today=date(2026, 4, 25), + max_lookback=20, + existing=set(), + ) + assert "2026-04-17" in out + + def test_results_sorted_chronologically(self): + out = _trading_days_to_process( + today=date(2026, 4, 25), + max_lookback=30, + existing=set(), + ) + assert out == sorted(out) + + def test_holiday_eval_date_skips_correctly(self): + # 2026-04-02 (Thu before Good Friday): trading day. fwd_5d via + # NYSE-aware arithmetic = 2026-04-10 (skip Good Friday). 4/10 < 4/25 + # → eligible, included. + out = _trading_days_to_process( + today=date(2026, 4, 25), + max_lookback=20, + existing=set(), + ) + assert "2026-04-02" in out