You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
811 lines
29 KiB
Python
811 lines
29 KiB
Python
"""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_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"
|