121 lines
5.0 KiB
Python
121 lines
5.0 KiB
Python
"""
|
|
Tests for collection job orchestration.
|
|
"""
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
from jobs.collection_job import run_daily_collection
|
|
|
|
|
|
def test_run_daily_collection_calls_collectors_in_order():
|
|
"""Daily collection should run all collectors in dependency order."""
|
|
call_order = []
|
|
|
|
def make_mock_collector(name):
|
|
mock_cls = MagicMock()
|
|
instance = MagicMock()
|
|
instance.run.side_effect = lambda: call_order.append(name)
|
|
mock_cls.return_value = instance
|
|
return mock_cls
|
|
|
|
with patch("jobs.collection_job.SessionLocal") as mock_session_local, \
|
|
patch("jobs.collection_job.StockCollector", make_mock_collector("stock")), \
|
|
patch("jobs.collection_job.SectorCollector", make_mock_collector("sector")), \
|
|
patch("jobs.collection_job.PriceCollector", make_mock_collector("price")), \
|
|
patch("jobs.collection_job.ValuationCollector", make_mock_collector("valuation")), \
|
|
patch("jobs.collection_job.ETFCollector", make_mock_collector("etf")), \
|
|
patch("jobs.collection_job.ETFPriceCollector", make_mock_collector("etf_price")):
|
|
mock_session_local.return_value = MagicMock()
|
|
run_daily_collection()
|
|
|
|
assert call_order == ["stock", "sector", "price", "valuation", "etf", "etf_price"]
|
|
|
|
|
|
def test_run_daily_collection_continues_on_failure():
|
|
"""If one collector fails, the rest should still run."""
|
|
call_order = []
|
|
|
|
def make_mock_collector(name, should_fail=False):
|
|
mock_cls = MagicMock()
|
|
instance = MagicMock()
|
|
def side_effect():
|
|
if should_fail:
|
|
raise RuntimeError(f"{name} failed")
|
|
call_order.append(name)
|
|
instance.run.side_effect = side_effect
|
|
mock_cls.return_value = instance
|
|
return mock_cls
|
|
|
|
with patch("jobs.collection_job.SessionLocal") as mock_session_local, \
|
|
patch("jobs.collection_job.StockCollector", make_mock_collector("stock", should_fail=True)), \
|
|
patch("jobs.collection_job.SectorCollector", make_mock_collector("sector")), \
|
|
patch("jobs.collection_job.PriceCollector", make_mock_collector("price")), \
|
|
patch("jobs.collection_job.ValuationCollector", make_mock_collector("valuation")), \
|
|
patch("jobs.collection_job.ETFCollector", make_mock_collector("etf")), \
|
|
patch("jobs.collection_job.ETFPriceCollector", make_mock_collector("etf_price")):
|
|
mock_session_local.return_value = MagicMock()
|
|
run_daily_collection()
|
|
|
|
# stock failed, but rest should continue
|
|
assert call_order == ["sector", "price", "valuation", "etf", "etf_price"]
|
|
|
|
|
|
from jobs.collection_job import run_backfill
|
|
|
|
|
|
def test_run_backfill_generates_yearly_chunks():
|
|
"""Backfill should split date range into yearly chunks."""
|
|
collected_ranges = []
|
|
|
|
def make_price_collector(name):
|
|
mock_cls = MagicMock()
|
|
def capture_init(db, start_date=None, end_date=None):
|
|
instance = MagicMock()
|
|
collected_ranges.append((name, start_date, end_date))
|
|
return instance
|
|
mock_cls.side_effect = capture_init
|
|
return mock_cls
|
|
|
|
with patch("jobs.collection_job.SessionLocal") as mock_session_local, \
|
|
patch("jobs.collection_job.PriceCollector", make_price_collector("price")), \
|
|
patch("jobs.collection_job.ETFPriceCollector", make_price_collector("etf_price")):
|
|
mock_db = MagicMock()
|
|
mock_session_local.return_value = mock_db
|
|
# Simulate no existing data (min date returns None)
|
|
mock_db.query.return_value.scalar.return_value = None
|
|
|
|
run_backfill(start_year=2023)
|
|
|
|
# Should generate chunks: 2023, 2024, 2025, 2026 (partial) for both price and etf_price
|
|
price_ranges = [(s, e) for name, s, e in collected_ranges if name == "price"]
|
|
assert len(price_ranges) >= 3 # At least 2023, 2024, 2025
|
|
assert price_ranges[0][0] == "20230101" # First chunk starts at start_year
|
|
|
|
|
|
def test_run_backfill_skips_already_collected_range():
|
|
"""Backfill should start from earliest existing data backwards."""
|
|
collected_ranges = []
|
|
|
|
def make_price_collector(name):
|
|
mock_cls = MagicMock()
|
|
def capture_init(db, start_date=None, end_date=None):
|
|
instance = MagicMock()
|
|
collected_ranges.append((name, start_date, end_date))
|
|
return instance
|
|
mock_cls.side_effect = capture_init
|
|
return mock_cls
|
|
|
|
from datetime import date
|
|
|
|
with patch("jobs.collection_job.SessionLocal") as mock_session_local, \
|
|
patch("jobs.collection_job.PriceCollector", make_price_collector("price")), \
|
|
patch("jobs.collection_job.ETFPriceCollector", make_price_collector("etf_price")):
|
|
mock_db = MagicMock()
|
|
mock_session_local.return_value = mock_db
|
|
|
|
# Simulate no existing data so both targets get backfilled
|
|
mock_db.query.return_value.scalar.return_value = None
|
|
# We'll verify the function runs without error
|
|
run_backfill(start_year=2023)
|
|
|
|
assert len(collected_ranges) > 0
|