"""Unit tests for snapshot export.""" from __future__ import annotations import json import tempfile import datetime as dt from pathlib import Path from unittest.mock import AsyncMock, MagicMock import pytest import pyarrow.parquet as pq from libs.export.snapshot_export import ( _enrich_prior_event_drift, _temporal_split, export_dataset_snapshot, ) @pytest.mark.unit class TestTemporalSplit: def test_split_proportions(self) -> None: """Temporal split produces correct proportions for 100 rows.""" rows = [{"event_date": "2026-01-01"} for _ in range(100)] splits = _temporal_split(rows, "temporal_70_15_15") assert len(splits["train"]) == 70 assert len(splits["valid"]) == 15 assert len(splits["test"]) == 15 def test_split_preserves_temporal_order(self) -> None: """Train set contains earliest dates, test contains latest.""" rows = [{"event_date": f"2026-{m:02d}-01"} for m in range(1, 13)] splits = _temporal_split(rows, "temporal_70_15_15") if splits["train"] and splits["test"]: assert splits["train"][-1]["event_date"] <= splits["test"][0]["event_date"] def test_empty_rows_returns_empty_splits(self) -> None: splits = _temporal_split([], "temporal_70_15_15") assert splits == {"train": [], "valid": [], "test": []} @pytest.mark.unit class TestPriorEventDrift: def test_prior_drift_only_populates_after_full_realization(self) -> None: rows = [ { "ticker": "ABC", "event_date": "2026-01-05", "entry_date": "2026-01-06", "fwd_return_5d": 0.12, }, { "ticker": "ABC", "event_date": "2026-01-14", "entry_date": "2026-01-15", "fwd_return_5d": 0.03, }, ] _enrich_prior_event_drift(rows) assert rows[0]["prior_event_fwd5d"] is None assert rows[1]["prior_event_fwd5d"] == pytest.approx(0.12) def test_prior_drift_stays_null_when_prior_window_not_finished(self) -> None: rows = [ { "ticker": "ABC", "event_date": "2026-01-05", "entry_date": "2026-01-06", "fwd_return_5d": 0.12, }, { "ticker": "ABC", "event_date": "2026-01-13", "entry_date": "2026-01-14", "fwd_return_5d": 0.03, }, ] _enrich_prior_event_drift(rows) assert rows[0]["prior_event_fwd5d"] is None assert rows[1]["prior_event_fwd5d"] is None @pytest.mark.unit class TestExportDatasetSnapshot: @pytest.mark.asyncio async def test_manifest_is_written(self) -> None: """export_dataset_snapshot writes a manifest.json with expected fields.""" # Mock DB session returning empty results (no feature+label pairs) mock_session = AsyncMock() mock_result = MagicMock() mock_result.all.return_value = [] # no rows mock_session.execute = AsyncMock(return_value=mock_result) with tempfile.TemporaryDirectory() as tmpdir: manifest = await export_dataset_snapshot( session=mock_session, snapshot_id="test-snapshot-001", split_policy="temporal_70_15_15", output_dir=tmpdir, ) assert manifest["snapshot_id"] == "test-snapshot-001" assert "created_at_utc" in manifest assert "row_counts" in manifest assert manifest["split_policy"] == "temporal_70_15_15" assert manifest["total_rows"] == 0 @pytest.mark.asyncio async def test_parquet_files_created(self) -> None: """Parquet files are created for each split partition.""" mock_session = AsyncMock() mock_result = MagicMock() mock_result.all.return_value = [] mock_session.execute = AsyncMock(return_value=mock_result) with tempfile.TemporaryDirectory() as tmpdir: await export_dataset_snapshot( session=mock_session, snapshot_id="test-parquet-002", split_policy="temporal_70_15_15", output_dir=tmpdir, ) snap_dir = Path(tmpdir) / "test-parquet-002" assert (snap_dir / "train.parquet").exists() assert (snap_dir / "valid.parquet").exists() assert (snap_dir / "test.parquet").exists() assert (snap_dir / "manifest.json").exists() manifest_data = json.loads((snap_dir / "manifest.json").read_text()) assert manifest_data["snapshot_id"] == "test-parquet-002" @pytest.mark.asyncio async def test_universe_profile_filters_and_enriches_rows(self, monkeypatch) -> None: mock_session = AsyncMock() mock_result = MagicMock() fs_pass = MagicMock(event_id="EVT::PASS", feature_json={ "reaction_day_return": 0.08, "volume_ratio_20d": 1.8, "avg_dollar_volume_20d": 90_000_000.0, "reaction_day_low": 95.0, "reaction_day_high": 105.0, "document_quality_score": 0.8, "signal_strength_score": 0.7, "guidance_direction_score": 1.0, "guidance_status": "raised", "oneoff_penalty": 0.1, "event_direction": "bullish", "parse_confidence_overall": 0.9, "parse_confidence_event_direction": 0.8, "parse_confidence_guidance": 0.75, }) fs_fail = MagicMock(event_id="EVT::FAIL", feature_json={ "reaction_day_return": 0.07, "volume_ratio_20d": 1.7, "avg_dollar_volume_20d": 10_000_000.0, }) lbl_pass = MagicMock( entry_convention="next_open", reaction_date=None, entry_date=None, entry_price=100.0, fwd_return_1d=0.01, fwd_return_3d=0.02, fwd_return_5d=0.03, hit_pos_1r_within_3d=False, hit_neg_1r_within_3d=False, close_up_after_3d=True, close_up_after_5d=True, mfe_3d=0.04, mae_3d=-0.01, mfe_5d=0.05, mae_5d=-0.02, fwd_return_10d=0.06, fwd_return_20d=0.07, mfe_10d=0.08, mae_10d=-0.03, mfe_20d=0.09, mae_20d=-0.04, label_status="ok", label_version="label-2.0.0", ) evt_pass = MagicMock(event_date=None) sym_pass = MagicMock(ticker="PASS", asset_type="common_stock", venue="NASDAQ") lbl_fail = MagicMock(**lbl_pass.__dict__) evt_fail = MagicMock(event_date=None) sym_fail = MagicMock(ticker="FAIL", asset_type="common_stock", venue="NASDAQ") mock_result.all.return_value = [ (fs_pass, lbl_pass, evt_pass, sym_pass), (fs_fail, lbl_fail, evt_fail, sym_fail), ] mock_session.execute = AsyncMock(return_value=mock_result) monkeypatch.setattr( "libs.export.snapshot_export._resolve_universe_profile", AsyncMock(return_value={ "PASS": {"market_cap_proxy": 5_000_000_000.0, "exchange_proxy": "NASDAQ"}, "FAIL": {"market_cap_proxy": 5_000_000_000.0, "exchange_proxy": "NASDAQ"}, }), ) with tempfile.TemporaryDirectory() as tmpdir: manifest = await export_dataset_snapshot( session=mock_session, snapshot_id="test-universe-profile", split_policy="temporal_70_15_15", output_dir=tmpdir, universe_profile="midlarge-liquid-long-v1", ) table = pq.read_table(Path(tmpdir) / "test-universe-profile" / "test.parquet") rows = table.to_pylist() assert manifest["universe_profile"] == "midlarge-liquid-long-v1" assert len(rows) == 1 assert rows[0]["market_cap_proxy"] == pytest.approx(5_000_000_000.0) assert rows[0]["reaction_day_low"] == pytest.approx(95.0) assert rows[0]["parse_confidence_guidance"] == pytest.approx(0.75) @pytest.mark.asyncio async def test_universe_profile_keeps_rows_when_export_enrichments_disabled(self, monkeypatch) -> None: mock_session = AsyncMock() mock_result = MagicMock() fs = MagicMock(event_id="EVT::PASS", feature_json={ "reaction_day_return": 0.08, "volume_ratio_20d": 1.8, }) lbl = MagicMock( entry_convention="next_open", reaction_date=None, entry_date=None, entry_price=100.0, fwd_return_1d=0.01, fwd_return_3d=0.02, fwd_return_5d=0.03, hit_pos_1r_within_3d=False, hit_neg_1r_within_3d=False, close_up_after_3d=True, close_up_after_5d=True, mfe_3d=0.04, mae_3d=-0.01, mfe_5d=0.05, mae_5d=-0.02, fwd_return_10d=0.06, fwd_return_20d=0.07, mfe_10d=0.08, mae_10d=-0.03, mfe_20d=0.09, mae_20d=-0.04, label_status="ok", label_version="label-2.0.0", ) evt = MagicMock(event_date=dt.date(2026, 1, 5), filed_at_utc=None, symbol_id="SYM1") sym = MagicMock(ticker="PASS", asset_type="common_stock", venue="NASDAQ") mock_result.all.return_value = [(fs, lbl, evt, sym)] mock_session.execute = AsyncMock(return_value=mock_result) monkeypatch.setattr( "libs.export.snapshot_export._resolve_universe_profile", AsyncMock(return_value={ "PASS": {"market_cap_proxy": 5_000_000_000.0, "exchange_proxy": "NASDAQ"}, }), ) async def _fake_backfill(rows): for row in rows: row["avg_dollar_volume_20d"] = 90_000_000.0 row["reaction_day_low"] = 95.0 row["reaction_day_high"] = 105.0 monkeypatch.setattr("libs.export.snapshot_export._backfill_market_fields", _fake_backfill) with tempfile.TemporaryDirectory() as tmpdir: manifest = await export_dataset_snapshot( session=mock_session, snapshot_id="test-universe-profile-no-export-enrich", split_policy="temporal_70_15_15", output_dir=tmpdir, universe_profile="midlarge-liquid-long-v1", include_export_enrichments=False, ) table = pq.read_table(Path(tmpdir) / "test-universe-profile-no-export-enrich" / "test.parquet") rows = table.to_pylist() assert manifest["total_rows"] == 1 assert manifest["export_enrichments"] == [] assert len(rows) == 1 assert rows[0]["avg_dollar_volume_20d"] == pytest.approx(90_000_000.0) @pytest.mark.asyncio async def test_midwide_universe_profile_uses_lower_thresholds(self, monkeypatch) -> None: mock_session = AsyncMock() mock_result = MagicMock() fs_midwide = MagicMock(event_id="EVT::MIDWIDE", feature_json={ "reaction_day_return": 0.06, "volume_ratio_20d": 1.4, "avg_dollar_volume_20d": 60_000_000.0, "reaction_day_low": 45.0, "reaction_day_high": 48.0, }) lbl = MagicMock( entry_convention="next_open", reaction_date=None, entry_date=None, entry_price=46.0, fwd_return_1d=0.01, fwd_return_3d=0.02, fwd_return_5d=0.03, hit_pos_1r_within_3d=False, hit_neg_1r_within_3d=False, close_up_after_3d=True, close_up_after_5d=True, mfe_3d=0.04, mae_3d=-0.01, mfe_5d=0.05, mae_5d=-0.02, fwd_return_10d=0.06, fwd_return_20d=0.07, mfe_10d=0.08, mae_10d=-0.03, mfe_20d=0.09, mae_20d=-0.04, label_status="ok", label_version="label-2.0.0", ) evt = MagicMock(event_date=None) sym = MagicMock(ticker="MIDW", asset_type="common_stock", venue="NASDAQ") mock_result.all.return_value = [(fs_midwide, lbl, evt, sym)] mock_session.execute = AsyncMock(return_value=mock_result) monkeypatch.setattr( "libs.export.snapshot_export._resolve_universe_profile", AsyncMock(return_value={ "MIDW": {"market_cap_proxy": 1_500_000_000.0, "exchange_proxy": "NASDAQ"}, }), ) with tempfile.TemporaryDirectory() as tmpdir: manifest = await export_dataset_snapshot( session=mock_session, snapshot_id="test-midwide-profile", split_policy="temporal_70_15_15", output_dir=tmpdir, universe_profile="midwide-liquid-long-v1", ) table = pq.read_table(Path(tmpdir) / "test-midwide-profile" / "test.parquet") rows = table.to_pylist() assert manifest["universe_profile"] == "midwide-liquid-long-v1" assert len(rows) == 1 assert rows[0]["market_cap_proxy"] == pytest.approx(1_500_000_000.0) assert rows[0]["avg_dollar_volume_20d"] == pytest.approx(60_000_000.0) @pytest.mark.asyncio async def test_default_export_merges_event_v1_with_market_v1(self) -> None: mock_session = AsyncMock() mock_result = MagicMock() fs_market = MagicMock( event_id="EVT::MERGE", feature_json={ "reaction_day_return": 0.08, "volume_ratio_20d": 2.1, }, ) fs_event = MagicMock( event_id="EVT::MERGE", feature_json={ "event_type": "earnings_release", "event_direction": "bullish", "guidance_status": "raised", "filing_time_bucket": "post_market", "parse_confidence_overall": 0.84, "document_quality_score": 0.79, }, ) lbl = MagicMock( entry_convention="next_open", reaction_date=None, entry_date=None, entry_price=100.0, fwd_return_1d=0.01, fwd_return_3d=0.02, fwd_return_5d=0.03, hit_pos_1r_within_3d=False, hit_neg_1r_within_3d=False, close_up_after_3d=True, close_up_after_5d=True, mfe_3d=0.04, mae_3d=-0.01, mfe_5d=0.05, mae_5d=-0.02, fwd_return_10d=0.06, fwd_return_20d=0.07, mfe_10d=0.08, mae_10d=-0.03, mfe_20d=0.09, mae_20d=-0.04, label_status="ok", label_version="label-2.0.0", invalid_event_for_labeling=False, ) evt = MagicMock(event_date=None) sym = MagicMock(ticker="MERG", asset_type="common_stock", venue="NASDAQ") mock_result.all.return_value = [ (fs_market, lbl, evt, sym), (fs_event, lbl, evt, sym), ] mock_session.execute = AsyncMock(return_value=mock_result) with tempfile.TemporaryDirectory() as tmpdir: manifest = await export_dataset_snapshot( session=mock_session, snapshot_id="test-merge-default", split_policy="temporal_70_15_15", output_dir=tmpdir, ) table = pq.read_table(Path(tmpdir) / "test-merge-default" / "test.parquet") rows = table.to_pylist() assert manifest["feature_version"] == "market_v1+event_v1" assert manifest["export_enrichments"] == [ "market_bar_backfill_v1", "pre_event_momentum_v1", "price_vs_sma20_v1", "prior_event_drift_v1", "macro_regime_v1", ] assert len(rows) == 1 assert rows[0]["event_type"] == "earnings_release" assert rows[0]["event_direction"] == "bullish" assert rows[0]["guidance_status"] == "raised" assert rows[0]["filing_time_bucket"] == "post_market" assert rows[0]["parse_confidence_overall"] == pytest.approx(0.84) assert rows[0]["document_quality_score"] == pytest.approx(0.79) @pytest.mark.asyncio async def test_export_preserves_multiple_entry_conventions_for_same_event(self) -> None: mock_session = AsyncMock() mock_result = MagicMock() fs_market = MagicMock( event_id="EVT::DUAL", snapshot_name="market_v1", feature_json={ "reaction_day_return": 0.08, "volume_ratio_20d": 2.1, "avg_dollar_volume_20d": 100_000_000.0, "reaction_day_low": 95.0, "reaction_day_high": 105.0, "pre_event_momentum_20d": 0.04, }, ) fs_event = MagicMock( event_id="EVT::DUAL", snapshot_name="event_v1", feature_json={ "event_type": "earnings_release", "event_direction": "bullish", "filing_time_bucket": "post_market", }, ) def make_label( label_id: int, entry_convention: str, entry_date: dt.date, entry_price: float, ) -> MagicMock: return MagicMock( label_id=label_id, entry_convention=entry_convention, reaction_date=dt.date(2026, 1, 5), entry_date=entry_date, entry_price=entry_price, fwd_return_1d=0.01, fwd_return_3d=0.02, fwd_return_5d=0.03, hit_pos_1r_within_3d=False, hit_neg_1r_within_3d=False, close_up_after_3d=True, close_up_after_5d=True, mfe_3d=0.04, mae_3d=-0.01, mfe_5d=0.05, mae_5d=-0.02, fwd_return_10d=0.06, fwd_return_20d=0.07, mfe_10d=0.08, mae_10d=-0.03, mfe_20d=0.09, mae_20d=-0.04, label_status="ok", label_version="label-2.0.0", ) lbl_next_open = make_label( 12, "next_open_after_reaction_close", dt.date(2026, 1, 6), 101.0, ) lbl_reaction_close = make_label( 11, "reaction_close", dt.date(2026, 1, 5), 100.0, ) evt = MagicMock(event_date=dt.date(2026, 1, 5), filed_at_utc=None) sym = MagicMock(ticker="DUAL", asset_type="common_stock", venue="NASDAQ") mock_result.all.return_value = [ (fs_event, lbl_next_open, evt, sym), (fs_market, lbl_reaction_close, evt, sym), (fs_market, lbl_next_open, evt, sym), (fs_event, lbl_reaction_close, evt, sym), ] mock_session.execute = AsyncMock(return_value=mock_result) with tempfile.TemporaryDirectory() as tmpdir: manifest = await export_dataset_snapshot( session=mock_session, snapshot_id="test-dual-entry-conventions", split_policy="temporal_100_0_0", output_dir=tmpdir, include_export_enrichments=False, ) table = pq.read_table(Path(tmpdir) / "test-dual-entry-conventions" / "train.parquet") rows = table.to_pylist() assert manifest["total_rows"] == 2 by_convention = {row["entry_convention"]: row for row in rows} assert set(by_convention) == { "reaction_close", "next_open_after_reaction_close", } assert by_convention["reaction_close"]["entry_date"] == "2026-01-05" assert by_convention["next_open_after_reaction_close"]["entry_date"] == "2026-01-06" assert by_convention["reaction_close"]["entry_price"] == pytest.approx(100.0) assert by_convention["next_open_after_reaction_close"]["entry_price"] == pytest.approx(101.0) assert by_convention["reaction_close"]["event_type"] == "earnings_release" assert by_convention["next_open_after_reaction_close"]["volume_ratio_20d"] == pytest.approx(2.1) @pytest.mark.asyncio async def test_export_falls_back_to_event_filed_time_bucket(self) -> None: mock_session = AsyncMock() mock_result = MagicMock() fs_market = MagicMock( event_id="EVT::FALLBACK", feature_json={ "reaction_day_return": 0.04, "volume_ratio_20d": 1.5, }, ) fs_event = MagicMock( event_id="EVT::FALLBACK", feature_json={ "event_type": "earnings_release", "event_direction": "bullish", "guidance_status": "raised", "filing_time_bucket": "unknown", "parse_confidence_overall": 0.8, }, ) lbl = MagicMock( entry_convention="next_open", reaction_date=None, entry_date=None, entry_price=100.0, fwd_return_1d=0.01, fwd_return_3d=0.02, fwd_return_5d=0.03, hit_pos_1r_within_3d=False, hit_neg_1r_within_3d=False, close_up_after_3d=True, close_up_after_5d=True, mfe_3d=0.04, mae_3d=-0.01, mfe_5d=0.05, mae_5d=-0.02, fwd_return_10d=0.06, fwd_return_20d=0.07, mfe_10d=0.08, mae_10d=-0.03, mfe_20d=0.09, mae_20d=-0.04, label_status="ok", label_version="label-2.0.0", invalid_event_for_labeling=False, ) evt = MagicMock( event_date=dt.date(2026, 1, 29), filed_at_utc=dt.datetime(2026, 1, 29, 22, 0, tzinfo=dt.UTC), ) sym = MagicMock(ticker="FBCK", asset_type="common_stock", venue="NASDAQ") mock_result.all.return_value = [ (fs_market, lbl, evt, sym), (fs_event, lbl, evt, sym), ] mock_session.execute = AsyncMock(return_value=mock_result) with tempfile.TemporaryDirectory() as tmpdir: await export_dataset_snapshot( session=mock_session, snapshot_id="test-filed-time-fallback", split_policy="temporal_70_15_15", output_dir=tmpdir, ) table = pq.read_table(Path(tmpdir) / "test-filed-time-fallback" / "test.parquet") rows = table.to_pylist() assert len(rows) == 1 assert rows[0]["filing_time_bucket"] == "post_market" @pytest.mark.asyncio async def test_symbol_exports_backfill_company_metadata_without_profile(self, monkeypatch) -> None: mock_session = AsyncMock() mock_result = MagicMock() fs_market = MagicMock( event_id="EVT::META", feature_json={ "reaction_day_return": 0.08, "volume_ratio_20d": 1.9, "avg_dollar_volume_20d": 90_000_000.0, "reaction_day_low": 95.0, "reaction_day_high": 105.0, }, ) lbl = MagicMock( entry_convention="next_open", reaction_date=None, entry_date=None, entry_price=100.0, fwd_return_1d=0.01, fwd_return_3d=0.02, fwd_return_5d=0.03, hit_pos_1r_within_3d=False, hit_neg_1r_within_3d=False, close_up_after_3d=True, close_up_after_5d=True, mfe_3d=0.04, mae_3d=-0.01, mfe_5d=0.05, mae_5d=-0.02, fwd_return_10d=0.06, fwd_return_20d=0.07, mfe_10d=0.08, mae_10d=-0.03, mfe_20d=0.09, mae_20d=-0.04, label_status="ok", label_version="label-2.0.0", invalid_event_for_labeling=False, ) evt = MagicMock(event_date=None) sym = MagicMock(ticker="META", asset_type="common_stock", venue="NASDAQ") mock_result.all.return_value = [(fs_market, lbl, evt, sym)] mock_session.execute = AsyncMock(return_value=mock_result) monkeypatch.setattr( "libs.export.snapshot_export._resolve_symbol_metadata", AsyncMock(return_value={ "META": {"market_cap_proxy": 6_500_000_000.0, "exchange_proxy": "NASDAQ"}, }), ) with tempfile.TemporaryDirectory() as tmpdir: await export_dataset_snapshot( session=mock_session, snapshot_id="test-symbol-metadata", split_policy="temporal_70_15_15", output_dir=tmpdir, symbols=["META"], ) table = pq.read_table(Path(tmpdir) / "test-symbol-metadata" / "test.parquet") rows = table.to_pylist() assert len(rows) == 1 assert rows[0]["market_cap_proxy"] == pytest.approx(6_500_000_000.0) assert rows[0]["exchange_proxy"] == "NASDAQ" @pytest.mark.asyncio async def test_midplus_universe_profile_requires_mid_liquidity(self, monkeypatch) -> None: mock_session = AsyncMock() mock_result = MagicMock() fs_pass = MagicMock(event_id="EVT::MIDPLUS::PASS", feature_json={ "reaction_day_return": 0.05, "volume_ratio_20d": 1.2, "avg_dollar_volume_20d": 70_000_000.0, "reaction_day_low": 31.0, "reaction_day_high": 35.0, }) fs_fail = MagicMock(event_id="EVT::MIDPLUS::FAIL", feature_json={ "reaction_day_return": 0.05, "volume_ratio_20d": 1.2, "avg_dollar_volume_20d": 60_000_000.0, "reaction_day_low": 31.0, "reaction_day_high": 35.0, }) lbl = MagicMock( entry_convention="next_open", reaction_date=None, entry_date=None, entry_price=32.0, fwd_return_1d=0.01, fwd_return_3d=0.02, fwd_return_5d=0.03, hit_pos_1r_within_3d=False, hit_neg_1r_within_3d=False, close_up_after_3d=True, close_up_after_5d=True, mfe_3d=0.04, mae_3d=-0.01, mfe_5d=0.05, mae_5d=-0.02, fwd_return_10d=0.06, fwd_return_20d=0.07, mfe_10d=0.08, mae_10d=-0.03, mfe_20d=0.09, mae_20d=-0.04, label_status="ok", label_version="label-2.0.0", ) evt = MagicMock(event_date=None) sym_pass = MagicMock(ticker="MPLS", asset_type="common_stock", venue="NASDAQ") sym_fail = MagicMock(ticker="MPLF", asset_type="common_stock", venue="NASDAQ") mock_result.all.return_value = [ (fs_pass, lbl, evt, sym_pass), (fs_fail, lbl, evt, sym_fail), ] mock_session.execute = AsyncMock(return_value=mock_result) monkeypatch.setattr( "libs.export.snapshot_export._resolve_universe_profile", AsyncMock(return_value={ "MPLS": {"market_cap_proxy": 2_000_000_000.0, "exchange_proxy": "NASDAQ"}, "MPLF": {"market_cap_proxy": 2_000_000_000.0, "exchange_proxy": "NASDAQ"}, }), ) with tempfile.TemporaryDirectory() as tmpdir: manifest = await export_dataset_snapshot( session=mock_session, snapshot_id="test-midplus-profile", split_policy="temporal_70_15_15", output_dir=tmpdir, universe_profile="midplus-liquid-long-v1", ) table = pq.read_table(Path(tmpdir) / "test-midplus-profile" / "test.parquet") rows = table.to_pylist() assert manifest["universe_profile"] == "midplus-liquid-long-v1" assert len(rows) == 1 assert rows[0]["avg_dollar_volume_20d"] == pytest.approx(70_000_000.0) @pytest.mark.asyncio async def test_universe_profile_falls_back_to_local_snapshot_metadata_when_screener_fails( self, monkeypatch, ) -> None: mock_session = AsyncMock() mock_result = MagicMock() fs_market = MagicMock( event_id="EVT::FALLBACK::PROFILE", feature_json={ "reaction_day_return": 0.08, "volume_ratio_20d": 1.9, "avg_dollar_volume_20d": 90_000_000.0, "reaction_day_low": 95.0, "reaction_day_high": 105.0, }, ) lbl = MagicMock( entry_convention="next_open", reaction_date=None, entry_date=None, entry_price=100.0, fwd_return_1d=0.01, fwd_return_3d=0.02, fwd_return_5d=0.03, hit_pos_1r_within_3d=False, hit_neg_1r_within_3d=False, close_up_after_3d=True, close_up_after_5d=True, mfe_3d=0.04, mae_3d=-0.01, mfe_5d=0.05, mae_5d=-0.02, fwd_return_10d=0.06, fwd_return_20d=0.07, mfe_10d=0.08, mae_10d=-0.03, mfe_20d=0.09, mae_20d=-0.04, label_status="ok", label_version="label-2.0.0", invalid_event_for_labeling=False, ) evt = MagicMock(event_date=None) sym = MagicMock(ticker="SAFE", asset_type="common_stock", venue="NASDAQ") mock_result.all.return_value = [(fs_market, lbl, evt, sym)] mock_session.execute = AsyncMock(return_value=mock_result) monkeypatch.setattr( "libs.export.snapshot_export._resolve_universe_profile", AsyncMock(side_effect=RuntimeError("oracle screener 500")), ) local_meta = { "SAFE": {"market_cap_proxy": 5_000_000_000.0, "exchange_proxy": "NASDAQ"}, } monkeypatch.setattr( "libs.export.snapshot_export._resolve_local_universe_profile_fallback", lambda universe_profile, exclude_snapshot_id=None: local_meta, ) symbol_meta = AsyncMock(return_value={}) monkeypatch.setattr("libs.export.snapshot_export._resolve_symbol_metadata", symbol_meta) with tempfile.TemporaryDirectory() as tmpdir: manifest = await export_dataset_snapshot( session=mock_session, snapshot_id="test-profile-fallback", split_policy="temporal_70_15_15", output_dir=tmpdir, universe_profile="midlarge-liquid-long-v1", ) table = pq.read_table(Path(tmpdir) / "test-profile-fallback" / "test.parquet") rows = table.to_pylist() symbol_meta.assert_not_awaited() assert manifest["universe_profile"] == "midlarge-liquid-long-v1" assert len(rows) == 1 assert rows[0]["market_cap_proxy"] == pytest.approx(5_000_000_000.0) assert rows[0]["exchange_proxy"] == "NASDAQ" @pytest.mark.asyncio async def test_date_range_filters_rows_and_persists_manifest(self) -> None: import datetime as dt mock_session = AsyncMock() mock_result = MagicMock() lbl = MagicMock( entry_convention="next_open", reaction_date=None, entry_date=None, entry_price=10.0, fwd_return_1d=0.0, fwd_return_3d=0.0, fwd_return_5d=0.0, hit_pos_1r_within_3d=False, hit_neg_1r_within_3d=False, close_up_after_3d=False, close_up_after_5d=False, mfe_3d=0.0, mae_3d=0.0, mfe_5d=0.0, mae_5d=0.0, fwd_return_10d=0.0, fwd_return_20d=0.0, mfe_10d=0.0, mae_10d=0.0, mfe_20d=0.0, mae_20d=0.0, label_status="ok", label_version="label-2.0.0", ) mock_result.all.return_value = [ ( MagicMock(event_id="EVT::A", feature_json={}), lbl, MagicMock(event_date=dt.date(2020, 6, 15)), MagicMock(ticker="AAA", asset_type="common_stock", venue="NASDAQ"), ), ( MagicMock(event_id="EVT::B", feature_json={}), lbl, MagicMock(event_date=dt.date(2022, 6, 15)), MagicMock(ticker="BBB", asset_type="common_stock", venue="NASDAQ"), ), ] mock_session.execute = AsyncMock(return_value=mock_result) with tempfile.TemporaryDirectory() as tmpdir: manifest = await export_dataset_snapshot( session=mock_session, snapshot_id="test-date-filter", split_policy="temporal_100_0_0", output_dir=tmpdir, start_date=dt.date(2020, 1, 1), end_date=dt.date(2020, 12, 31), ) table = pq.read_table(Path(tmpdir) / "test-date-filter" / "train.parquet") rows = table.to_pylist() assert manifest["start_date"] == "2020-01-01" assert manifest["end_date"] == "2020-12-31" assert len(rows) == 1 assert rows[0]["event_id"] == "EVT::A"