Mock jretrieve

ci/cscs.yml

@@ -11,33 +11,19 @@ include:
     SLURM_JOB_NUM_NODES: 1
     SLURM_NTASKS: 1
 
-unit_test_job:
+integration_test_job:
   extends: [.baremetal-runner-balfrin]
   before_script:
     - |
-      echo "Setting up environment for unit tests on $(hostname)"
+      echo "Setting up environment for integration tests on $(hostname)"
       VENV_DIR="$(pwd)/.venv"
       echo "uv installed at: $(which uv)"
       uv --version
       echo "will install evalml at: $VENV_DIR"
       UV_HOME="$VENV_DIR" uv sync
       source "$VENV_DIR/bin/activate"
   script:
-    - pytest tests/unit
+    - pytest tests/integration -m longtest
   variables:
     # Explicitly request more time from Slurm, tests aren't running fast enough.
     SLURM_TIMELIMIT: "01:00:00"
-
-integration_test_job:
-  extends: [.baremetal-runner-balfrin]
-  before_script:
-    - |
-      echo "Setting up environment for integration tests on $(hostname)"
-      VENV_DIR="$(pwd)/.venv"
-      echo "uv installed at: $(which uv)"
-      uv --version
-      echo "will install evalml at: $VENV_DIR"
-      UV_HOME="$VENV_DIR" uv sync
-      source "$VENV_DIR/bin/activate"
-  script:
-    - pytest tests/integration

tests/integration/capture_jretrieve_fixtures.py

@@ -0,0 +1,58 @@
+"""Capture real jretrieve output and write it to fixture CSV files.
+
+Run once locally with DWH credentials available to refresh the fixtures used by
+the mock_jretrieve pytest fixture:
+
+    python tests/integration/capture_jretrieve_fixtures.py
+
+The script intercepts the raw CSV strings returned by _run_with_retry and
+writes them to tests/integration/fixtures/jretrieve/{meta,data}.csv.
+"""
+
+from datetime import datetime
+from pathlib import Path
+from unittest.mock import patch
+
+import sys
+
+sys.path.insert(0, str(Path(__file__).parents[2] / "src"))
+
+from data_input import jretrieve as jr
+import data_input
+
+FIXTURE_DIR = Path(__file__).parent / "fixtures" / "jretrieve"
+ROOT = "jretrievedwh:1,2"
+REFTIME = datetime(2024, 8, 1, 0, 0)
+STEPS = [0, 6, 12]
+PARAMS = ["T_2M", "SP_10M", "TOT_PREC6"]
+
+captured: dict[str, str] = {}
+_original_run_with_retry = jr._run_with_retry
+
+
+def _capturing_run_with_retry(argv, env, timeout_s, attempts=3):
+    result = _original_run_with_retry(
+        argv, env=env, timeout_s=timeout_s, attempts=attempts
+    )
+    call_type = "meta" if "--meta-info" in argv else "data"
+    captured[call_type] = result
+    return result
+
+
+def main():
+    FIXTURE_DIR.mkdir(parents=True, exist_ok=True)
+    with patch.object(jr, "_run_with_retry", _capturing_run_with_retry):
+        data_input.load_obs_data_from_jretrieve(ROOT, REFTIME, STEPS, PARAMS)
+
+    for call_type, csv_text in captured.items():
+        out = FIXTURE_DIR / f"{call_type}.csv"
+        out.write_text(csv_text)
+        print(f"written: {out} ({len(csv_text.splitlines())} lines)")
+
+    if len(captured) < 2:
+        missing = {"meta", "data"} - captured.keys()
+        print(f"WARNING: missing captures for: {missing}")
+
+
+if __name__ == "__main__":
+    main()

tests/integration/configs/meteogram.yaml → tests/integration/configs/meteogram_small.yaml

@@ -6,14 +6,15 @@ description: |
 config_label: meteogram-test
 
 dates:
-  - 2024-01-01T00:00
+  - 2024-08-01T00:00
 
 runs:
   - forecaster:
       checkpoint: https://servicedepl.meteoswiss.ch/mlstore#/experiments/409/runs/b30acf68520a4bbd8324c44666561696
       label: stage_C_icon_1km
-      steps: 0/6/6
+      steps: 0/12/6
       config: resources/inference/configs/sgm-forecaster-global-ich1.yaml
+      squash_venv: false
       extra_requirements:
         - earthkit-utils<0.2.0
         - earthkit-data<0.19.0
@@ -43,13 +44,17 @@ showcase:
   params:
     - T_2M
     - SP_10M
+    - TOT_PREC6
   meteograms:
     enabled: true
     stations:
       - GVE
+      - SAE
   animations:
     enabled: false
 
+lapse_rate_correction: false
+
 locations:
   output_root: output/
 

tests/integration/conftest.py

@@ -0,0 +1,24 @@
+import os
+from pathlib import Path
+
+import pytest
+
+_FIXTURE_DIR = Path(__file__).parent / "fixtures" / "jretrieve"
+_BIN_DIR = _FIXTURE_DIR / "bin"
+
+
+@pytest.fixture
+def mock_jretrieve(monkeypatch):
+    """Intercept jretrieve by placing the mock binary first on PATH.
+
+    fixtures/jretrieve/bin/jretrievedwh.py lives on the shared /scratch
+    filesystem, so it is accessible from both the pytest node and any SLURM
+    compute nodes that Snakemake dispatches plot_meteogram jobs to.
+
+    monkeypatch restores os.environ to its original state after the test, so
+    real credentials and PATH are unaffected in subsequent runs.
+    """
+    monkeypatch.setenv("PATH", f"{_BIN_DIR}:{os.environ.get('PATH', '')}")
+    monkeypatch.setenv("JRETRIEVE_FIXTURE_DIR", str(_FIXTURE_DIR))
+    monkeypatch.setenv("JRETRIEVE_CLIENT_ID", "mock")
+    monkeypatch.setenv("JRETRIEVE_CLIENT_SECRET", "mock")

tests/integration/fixtures/jretrieve/bin/jretrievedwh.py

@@ -0,0 +1,14 @@
+#!/usr/bin/env python3
+"""Mock jretrievedwh.py binary for integration tests.
+
+Returns pre-stored fixture CSVs instead of contacting the DWH.
+The fixture directory is read from JRETRIEVE_FIXTURE_DIR.
+"""
+
+import os
+import sys
+from pathlib import Path
+
+fixture_dir = Path(os.environ["JRETRIEVE_FIXTURE_DIR"])
+call_type = "meta" if "--meta-info" in sys.argv else "data"
+sys.stdout.write((fixture_dir / f"{call_type}.csv").read_text())