wehub-resource-sync
157 lines
4.9 KiB
Python
157 lines
4.9 KiB
Python
"""Test case study pipeline notebooks via Papermill parameter injection.
|
|
|
|
Each case study notebook runs independently against pre-generated intermediates
|
|
(labels, features, predictions, registries) stored in the test-data repo.
|
|
A failure in one notebook does NOT cascade to skip later notebooks.
|
|
|
|
Stages are auto-discovered: any [0-9][0-9]_*.py file in a case study
|
|
directory is treated as a pipeline stage.
|
|
|
|
Usage:
|
|
# All case studies
|
|
pytest tests/test_case_studies.py -v
|
|
|
|
# Specific case study
|
|
pytest tests/test_case_studies.py -v -k "etfs"
|
|
|
|
# Specific stage
|
|
pytest tests/test_case_studies.py -v -k "03_features"
|
|
"""
|
|
|
|
import re
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from tests.pm_helpers import current_test_tier, get_overrides, get_tier, run_notebook
|
|
|
|
REPO_ROOT = Path(__file__).parent.parent
|
|
|
|
# All case studies
|
|
CASE_STUDIES = [
|
|
"etfs",
|
|
"crypto_perps_funding",
|
|
"nasdaq100_microstructure",
|
|
"sp500_equity_option_analytics",
|
|
"us_firm_characteristics",
|
|
"fx_pairs",
|
|
"cme_futures",
|
|
"sp500_options",
|
|
"us_equities_panel",
|
|
]
|
|
|
|
# Pattern for numbered pipeline stages — allows optional single-letter suffix
|
|
# (e.g., 10a_pca, 11b_ipca) for per-estimator notebook splits.
|
|
_STAGE_RE = re.compile(r"^\d{2}[a-z]?_")
|
|
|
|
|
|
def _collect_case_study_tests():
|
|
"""Collect all case study pipeline notebooks as (case_study, stage, path) tuples.
|
|
|
|
Auto-discovers files matching ^\\d{2}[a-z]?_ in each case study directory,
|
|
sorted numerically. Skips helper files (starting with _).
|
|
"""
|
|
tests = []
|
|
for cs in CASE_STUDIES:
|
|
cs_dir = REPO_ROOT / "case_studies" / cs
|
|
if not cs_dir.exists():
|
|
continue
|
|
|
|
for notebook in sorted(cs_dir.glob("[0-9][0-9]*.py")):
|
|
if notebook.name.startswith("_"):
|
|
continue
|
|
if not _STAGE_RE.match(notebook.name):
|
|
continue
|
|
stage = notebook.stem # e.g., "06_linear" or "11a_pca"
|
|
tests.append((cs, stage, notebook))
|
|
|
|
return tests
|
|
|
|
|
|
CASE_STUDY_TESTS = _collect_case_study_tests()
|
|
|
|
print(f"Found {len(CASE_STUDY_TESTS)} case study pipeline notebooks to test")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"case_study,stage,notebook_path",
|
|
CASE_STUDY_TESTS,
|
|
ids=lambda *args: None, # Custom IDs below
|
|
)
|
|
def test_case_study_pipeline(
|
|
case_study, stage, notebook_path, populated_data_dir, seeded_output_dir
|
|
):
|
|
"""Execute a case study pipeline stage via Papermill.
|
|
|
|
Each notebook runs independently — intermediates (labels, features,
|
|
predictions, registries) are pre-generated in the test-data repo.
|
|
"""
|
|
# Check case-study-level skip (e.g., "case_studies/nasdaq100_microstructure")
|
|
cs_key = f"case_studies/{case_study}"
|
|
cs_overrides = get_overrides(cs_key)
|
|
if cs_overrides.get("skip"):
|
|
pytest.skip(f"Skipped: {cs_overrides.get('skip_reason', 'case study skipped')}")
|
|
|
|
rel_path = notebook_path.relative_to(REPO_ROOT).with_suffix("")
|
|
overrides = get_overrides(str(rel_path))
|
|
|
|
# Tier routing: skip when NB tier doesn't match the current run tier.
|
|
nb_tier = get_tier(overrides)
|
|
run_tier = current_test_tier()
|
|
if nb_tier != run_tier:
|
|
pytest.skip(f"Tier {nb_tier} — current run tier is {run_tier}")
|
|
|
|
# Skip if overrides say so
|
|
if overrides.get("skip"):
|
|
reason = overrides.get("skip_reason", "marked skip in overrides")
|
|
pytest.skip(f"Skipped: {reason}")
|
|
|
|
# Check required imports (e.g., gensim, signatory, duckdb)
|
|
requires = overrides.get("requires_import")
|
|
if requires:
|
|
pkg = requires if isinstance(requires, str) else requires[0]
|
|
try:
|
|
__import__(pkg)
|
|
except ImportError:
|
|
pytest.skip(f"Requires {pkg} (not installed in this Docker image)")
|
|
|
|
# Check GPU requirement
|
|
if overrides.get("gpu"):
|
|
try:
|
|
import torch
|
|
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("GPU required but not available")
|
|
except ImportError:
|
|
pytest.skip("GPU required but torch not installed")
|
|
|
|
timeout = overrides.get("timeout", 300)
|
|
parameters = overrides.get("parameters", {})
|
|
|
|
result = run_notebook(
|
|
py_path=notebook_path,
|
|
parameters=parameters,
|
|
timeout=timeout,
|
|
output_dir=seeded_output_dir,
|
|
data_dir=populated_data_dir,
|
|
)
|
|
|
|
if result["status"] == "error":
|
|
pytest.fail(
|
|
f"\n{'=' * 70}\n"
|
|
f"Pipeline failed: {case_study}::{stage}\n"
|
|
f"{'=' * 70}\n"
|
|
f"Error: {result['error']}\n"
|
|
f"{'=' * 70}\n"
|
|
)
|
|
|
|
|
|
# Custom test IDs
|
|
def pytest_collection_modifyitems(items):
|
|
"""Set readable test IDs for case study tests."""
|
|
for item in items:
|
|
if "test_case_study_pipeline" in item.name and hasattr(item, "callspec"):
|
|
cs = item.callspec.params.get("case_study", "")
|
|
stage = item.callspec.params.get("stage", "")
|
|
item._nodeid = f"{item.parent.nodeid}::{cs}::{stage}"
|