fix tests, update CI requirement (#14)

* fix failing tests

* add ci helper

* ruff

* fix false positives

* fix import error

* fix CI trigger: add pull_request event so checks report to PR

workflow_dispatch runs don't appear in the PR Checks tab — only
pull_request-triggered runs are associated with PRs. Add
pull_request: [opened, reopened] so the check runs once when a PR
is opened (not on every push). Keep workflow_dispatch for manual
re-runs.

Also fix branch protection context name: the check run is named
'test' (job name), not 'test / test'.

Author: jeffbrennan

.github/workflows/test.yml

@@ -1,10 +1,26 @@
 name: test
 
 on:
-  push:
-    branches: ["main"]
   pull_request:
-    branches: ["main"]
+    types: [opened, reopened]
+  workflow_dispatch:
+
+# Runs automatically when a PR is opened or reopened, but NOT on every push.
+# To re-run after pushing new commits, use "Re-run all checks" in the PR's
+# Checks tab, or trigger manually from Actions → test → Run workflow.
+#
+# Branch protection on `main` requires the `test` check to pass before merge:
+#   gh api -X PUT repos/jeffbrennan/sparkparse/branches/main/protection --input - <<'EOF'
+#   {
+#     "required_status_checks": {
+#       "strict": false,
+#       "contexts": ["test"]
+#     },
+#     "enforce_admins": false,
+#     "restrictions": null,
+#     "required_pull_request_reviews": null
+#   }
+#   EOF
 
 jobs:
   test:
@@ -30,7 +46,7 @@ jobs:
         run: uv run ruff format --check sparkparse/ tests/
 
       - name: Type check
-        run: uv run pyrefly check sparkparse/
+        run: uv run pyrefly check sparkparse/ tests/
 
       - name: Unit tests
         run: uv run pytest tests/ --ignore=tests/test.py --ignore=tests/test_capture.py -v

CLAUDE.md

@@ -31,24 +31,26 @@ Spark event log (JSONL)
 
 ## Key files
 
-| File | Purpose |
-|---|---|
-| `sparkparse/models.py` | All Pydantic models. `NodeType` enum (30+ values), `ParsedLog`, `ParsedLogDataFrames`, `NODE_TYPE_DETAIL_MAP` (node type → detail model class) |
-| `sparkparse/parse.py` | `get_parsed_metrics()` is the main entry point. `parse_spark_ui_tree()` converts indented ASCII plans to node graphs. `parse_log()` orchestrates everything. |
-| `sparkparse/clean.py` | `log_to_dag_df()` and `log_to_combined_df()` produce the two output DataFrames. `get_readable_size()` and `get_readable_timing()` are Polars expression helpers. |
-| `sparkparse/app.py` | Typer CLI. `get` → parses and writes output files. `viz` → launches dashboard. |
-| `sparkparse/capture.py` | `SparkparseCapture` context manager/decorator. Stops the active session, restarts it with event logging enabled, then processes logs on `__exit__`. |
+| File                    | Purpose                                                                                                                                                          |
+| ----------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `sparkparse/models.py`  | All Pydantic models. `NodeType` enum (30+ values), `ParsedLog`, `ParsedLogDataFrames`, `NODE_TYPE_DETAIL_MAP` (node type → detail model class)                   |
+| `sparkparse/parse.py`   | `get_parsed_metrics()` is the main entry point. `parse_spark_ui_tree()` converts indented ASCII plans to node graphs. `parse_log()` orchestrates everything.     |
+| `sparkparse/clean.py`   | `log_to_dag_df()` and `log_to_combined_df()` produce the two output DataFrames. `get_readable_size()` and `get_readable_timing()` are Polars expression helpers. |
+| `sparkparse/app.py`     | Typer CLI. `get` → parses and writes output files. `viz` → launches dashboard.                                                                                   |
+| `sparkparse/capture.py` | `SparkparseCapture` context manager/decorator. Stops the active session, restarts it with event logging enabled, then processes logs on `__exit__`.              |
 
 ## Data model
 
 ### `dag` DataFrame columns (per physical plan node per query)
+
 - `query_id`, `query_function`, `query_header`, `query_start/end_timestamp`, `query_duration_seconds`
 - `node_id`, `node_type`, `node_name`, `child_nodes` (comma-separated string), `whole_stage_codegen_id`
 - `details` — JSON string; deserialize with the appropriate model from `NODE_TYPE_DETAIL_MAP`
 - `accumulator_totals` — list of structs: `{metric_name, metric_type, value, readable_str, unit}`
 - `node_duration_minutes`
 
 ### `combined` DataFrame columns (per task)
+
 - `log_name`, `parsed_log_name`, `query_id`, `query_function`
 - `query/job/stage/task` start/end timestamps and duration_seconds
 - `task_id`, `executor_id`, `nodes` (list of physical plan nodes for this task)
@@ -84,6 +86,7 @@ uv run pytest tests/test.py tests/test_capture.py -v
 ```
 
 Test fixtures live in `tests/data/`:
+
 - `tests/data/full_logs/` — three real Spark event logs for end-to-end testing
 - `tests/data/test_*_parsing/` — extracted JSON/TXT fragments for unit testing specific parsing steps
 
@@ -93,57 +96,63 @@ Test fixtures live in `tests/data/`:
 uv sync --dev                          # install including dev deps
 uv run ruff check sparkparse/ tests/   # lint
 uv run ruff format sparkparse/ tests/  # format
-uv run pyrefly check sparkparse/       # type check
+uv run pyrefly check sparkparse/ tests/  # type check
 ```
 
 ## Planned improvements (see IMPLEMENTATION.md for full detail)
 
 ### PR 1 — Tooling modernization
+
 - Move dev deps (`pytest`, `ruff`, `ipykernel`) to `[tool.uv.dev-dependencies]`
 - Add `pyrefly` for type checking
 - Add `[tool.ruff]` config to `pyproject.toml`
 - Add `.github/workflows/test.yml` CI workflow
 
 ### PR 2 — Unit test expansion
+
 - Add `tests/test_clean.py` covering `get_readable_size`, `get_readable_timing`, `clean_jobs`,
   `clean_stages`, `clean_tasks`, `get_job_idle_time`
 
 ### PR 3 — LLM-friendly analysis output
+
 - Add `sparkparse/analyze.py` with two distinct concerns:
-  - `to_plan_summary(dfs, log_name) -> dict` — token-efficient structured plan data for LLM
-    piping; presents raw facts (nodes, durations, bytes, join types, paths) without pre-assigned
-    severity or pre-classified findings; the LLM draws its own conclusions. The value over
-    `df.explain()` is runtime metrics (per-node durations, scan bytes/records) correlated from
-    accumulator updates — plan structure alone is not worth re-encoding.
-  - `find_*()` helpers — programmatic analysis functions for interactive/notebook use
-    (`find_cartesian_joins`, `find_largest_scans`, `find_repeated_scans`, `find_spill`, etc.)
+    - `to_plan_summary(dfs, log_name) -> dict` — token-efficient structured plan data for LLM
+      piping; presents raw facts (nodes, durations, bytes, join types, paths) without pre-assigned
+      severity or pre-classified findings; the LLM draws its own conclusions. The value over
+      `df.explain()` is runtime metrics (per-node durations, scan bytes/records) correlated from
+      accumulator updates — plan structure alone is not worth re-encoding.
+    - `find_*()` helpers — programmatic analysis functions for interactive/notebook use
+      (`find_cartesian_joins`, `find_largest_scans`, `find_repeated_scans`, `find_spill`, etc.)
 - Add `tests/test_analyze.py`
 
 ### PR 4 — CLI improvements
+
 - Add `sparkparse analyze` command (prints JSON findings, pipeable to LLM)
 - Improve help strings on existing commands
 - Add `"analyze"` action to `capture.py`
 - Replace `print()` with `logging` in `capture.py`
 
 ### PR 7 — Complete Pydantic coverage
+
 - Add `NodeType.Unknown` sentinel and `RawDetail` fallback model so unrecognized node types
   degrade gracefully instead of raising (currently `parse_spark_ui_tree` and `get_plan_details`
   both hard-fail on unknown types — see `parse.py:278` and `parse.py:145`)
 - Wrap `NodeType(...)` and `QueryFunction(...)` calls in try/except with warning logs
 - Add detail models for the standard Spark 3.5 operator set not yet covered:
-  - Aggregate variants: `ObjectHashAggregate`, `SortAggregate`, `Expand`
-  - Python/Pandas UDF nodes: `ArrowEvalPython`, `BatchEvalPython`, `MapInPandas`,
-    `MapInArrow`, `FlatMapGroupsInPandas`, `FlatMapCoGroupsInPandas`
-  - DataSource V2: `BatchScanExec`, `WriteToDatasourceV2`, `AppendData`,
-    `OverwriteByExpression`, `OverwritePartitionsDynamic`
-  - Subquery: `SubqueryExec`, `ReusedSubqueryExec`, `SubqueryBroadcast`
-  - Misc: `RepartitionByExpression`, `Sample`, `Range`, `CartesianProduct`
+    - Aggregate variants: `ObjectHashAggregate`, `SortAggregate`, `Expand`
+    - Python/Pandas UDF nodes: `ArrowEvalPython`, `BatchEvalPython`, `MapInPandas`,
+      `MapInArrow`, `FlatMapGroupsInPandas`, `FlatMapCoGroupsInPandas`
+    - DataSource V2: `BatchScanExec`, `WriteToDatasourceV2`, `AppendData`,
+      `OverwriteByExpression`, `OverwritePartitionsDynamic`
+    - Subquery: `SubqueryExec`, `ReusedSubqueryExec`, `SubqueryBroadcast`
+    - Misc: `RepartitionByExpression`, `Sample`, `Range`, `CartesianProduct`
 - Extend `ExchangeType` with `REPARTITION_BY_COL`, `REPARTITION_BY_NUM`, `REPARTITION`
 - Extend `QueryFunction` with `show`, `collect`, `first`, `head`, `take`, etc.
 - Databricks/Photon-specific nodes handled by `Unknown` fallback until real fixtures available
 - Add test fixtures and cases for each new node type in `tests/test_detail_parsing.py`
 
 ### PR 5 — Cloud storage support
+
 - Add `sparkparse/storage.py` with path-agnostic I/O (`open_file`, `write_text`, `list_files`,
   `copy_file`, `remove_dir`) backed by `fsspec` for cloud URIs and stdlib for local paths
 - Update `parse.py`, `common.py`, `capture.py` to route all file I/O through `storage.py`
@@ -154,6 +163,7 @@ uv run pyrefly check sparkparse/       # type check
   cluster termination
 
 ### PR 6 — Performance history and alerts
+
 - Add `sparkparse/history.py` — append-only run history: `record_from_dfs()` derives a compact
   `RunRecord` (duration, bytes, spill, shuffle, cartesian join count, etc.) from
   `ParsedLogDataFrames`; `append()` writes to Delta (primary) or JSONL (fallback); `read()`

Justfile

@@ -1,9 +1,22 @@
 ci:
     uv run ruff check sparkparse/ tests/
     uv run ruff format --check sparkparse/ tests/
-    uv run pyrefly check sparkparse/
+    uv run pyrefly check sparkparse/ tests/
     uv run pytest tests/ --ignore=tests/test.py --ignore=tests/test_capture.py -v
 
+# Trigger the test workflow on the current branch and watch it run.
+# The workflow also auto-runs when a PR is opened/reopened; use this for
+# manual re-runs after pushing new commits. Requires gh auth login.
+run-ci:
+    @branch=$(git branch --show-current); \
+    echo "Triggering test workflow on branch: $branch"; \
+    gh workflow run test.yml --ref $branch || { echo "Failed to trigger workflow. Is the branch pushed to GitHub?"; exit 1; }; \
+    sleep 3; \
+    run_id=$(gh run list --workflow=test.yml --branch=$branch --limit=1 --json databaseId --jq '.[0].databaseId'); \
+    if [ -z "$run_id" ]; then echo "Could not find the triggered run. Check: gh run list --workflow=test.yml"; exit 1; fi; \
+    echo "Watching run $run_id..."; \
+    gh run watch $run_id --exit-status && echo "CI passed" || { echo "CI failed"; gh run view $run_id --log-failed; exit 1; }
+
 jsonformat:
     mkdir -p data/logs/sandbox
     for file in data/logs/raw/*; do \
@@ -15,4 +28,4 @@ live:
     open http://localhost:4040/
 
 history:
-    ($SPARK_HOME/sbin/start-history-server.sh || open http://localhost:18080/)
+    ($SPARK_HOME/sbin/start-history-server.sh || open http://localhost:18080/)

pyproject.toml

@@ -46,6 +46,7 @@ dev-dependencies = [
 [tool.ruff]
 target-version = "py311"
 line-length = 88
+exclude = ["typings/**"]
 
 [tool.ruff.lint]
 select = ["E", "F", "I", "UP"]
@@ -59,3 +60,6 @@ project-includes = [
     "**/*.py*",
     "**/*.ipynb",
 ]
+project-excludes = ["typings/**"]
+search-path = ["typings"]
+ignore-missing-imports = ["deltalake"]

sparkparse/alerts.py

@@ -77,7 +77,9 @@ def load_alert_config(path: str) -> list[AlertConfig]:
     return [AlertConfig(**a) for a in raw_alerts]
 
 
-def _compute_baseline(history: pl.DataFrame, alert: AlertConfig, current_run_id: str) -> float:
+def _compute_baseline(
+    history: pl.DataFrame, alert: AlertConfig, current_run_id: str
+) -> float:
     """Mean of the alert metric over the last ``window`` runs, excluding the
     current run. Returns 0.0 when no history is available.
     """
@@ -86,7 +88,8 @@ def _compute_baseline(history: pl.DataFrame, alert: AlertConfig, current_run_id:
 
     hist = (
         history.filter(
-            (pl.col("log_name") == alert.log_name) & (pl.col("run_id") != current_run_id)
+            (pl.col("log_name") == alert.log_name)
+            & (pl.col("run_id") != current_run_id)
         )
         .sort("run_at", descending=True)
         .head(alert.window)
@@ -96,10 +99,14 @@ def _compute_baseline(history: pl.DataFrame, alert: AlertConfig, current_run_id:
         return 0.0
 
     mean_val = hist[alert.metric].mean()
-    return float(mean_val) if mean_val is not None else 0.0
+    if mean_val is not None and isinstance(mean_val, int | float):
+        return float(mean_val)
+    return 0.0
 
 
-def _dispatch(alert: AlertConfig, alert_dict: dict, alert_output_path: str | None) -> None:
+def _dispatch(
+    alert: AlertConfig, alert_dict: dict, alert_output_path: str | None
+) -> None:
     if alert.on_trigger == "raise":
         raise SparkparseAlertError(
             alert_name=alert.name,

sparkparse/analyze.py

@@ -199,12 +199,16 @@ def find_largest_scans(dfs: ParsedLogDataFrames, n: int = 10) -> pl.DataFrame:
     )
 
     return (
-        scan_nodes.select("query_id", "node_id", "node_name", "details", "node_duration_minutes")
+        scan_nodes.select(
+            "query_id", "node_id", "node_name", "details", "node_duration_minutes"
+        )
         .join(node_bytes, on=["query_id", "node_name"], how="left")
         .with_columns(
             pl.col("details")
             .map_elements(
-                lambda s: json.loads(s)["detail"]["location"]["location"] if s is not None else [],
+                lambda s: json.loads(s)["detail"]["location"]["location"]
+                if s is not None
+                else [],
                 return_dtype=pl.List(pl.String),
             )
             .alias("paths")

sparkparse/app.py

@@ -111,7 +111,9 @@ def analyze(
     ] = None,
     out_file: Annotated[
         str | None,
-        typer.Option(help="Write analysis output to this file (local path or cloud URI)."),
+        typer.Option(
+            help="Write analysis output to this file (local path or cloud URI)."
+        ),
     ] = None,
     format: Annotated[
         AnalysisFormat,
@@ -140,7 +142,9 @@ def analyze(
         lines.append(f"Bytes read: {totals.get('bytes_read', 0):,}")
         lines.append(f"Bytes written: {totals.get('bytes_written', 0):,}")
         lines.append(f"Shuffle bytes read: {totals.get('shuffle_bytes_read', 0):,}")
-        lines.append(f"Shuffle bytes written: {totals.get('shuffle_bytes_written', 0):,}")
+        lines.append(
+            f"Shuffle bytes written: {totals.get('shuffle_bytes_written', 0):,}"
+        )
         lines.append(f"Memory spilled: {totals.get('memory_bytes_spilled', 0):,}")
         lines.append(f"Disk spilled: {totals.get('disk_bytes_spilled', 0):,}")
         for q in summary.get("queries", []):
@@ -202,11 +206,17 @@ def check_alerts_cmd(
     history_path: Annotated[
         str, typer.Argument(help="Path to history store (Delta dir or JSONL file).")
     ],
-    log_name: Annotated[str, typer.Argument(help="Stable job identifier to check alerts for.")],
-    alert_config: Annotated[str, typer.Argument(help="Path to alert configuration TOML file.")],
+    log_name: Annotated[
+        str, typer.Argument(help="Stable job identifier to check alerts for.")
+    ],
+    alert_config: Annotated[
+        str, typer.Argument(help="Path to alert configuration TOML file.")
+    ],
     alert_output_path: Annotated[
         str | None,
-        typer.Option(help="File to write triggered alerts (for on_trigger='file' rules)."),
+        typer.Option(
+            help="File to write triggered alerts (for on_trigger='file' rules)."
+        ),
     ] = None,
 ) -> None:
     """Run alert checks against the latest run in history."""

sparkparse/capture.py

@@ -63,7 +63,9 @@ def __init__(
         self._last_record: RunRecord | None = None
         self._triggered_alerts: list[dict] = []
 
-    def __call__(self, func: Callable[..., R]) -> Callable[..., tuple[R, "SparkparseCapture"]]:
+    def __call__(
+        self, func: Callable[..., R]
+    ) -> Callable[..., tuple[R, "SparkparseCapture"]]:
         @functools.wraps(func)
         def get_wrapper(*args: Any, **kwargs: Any) -> tuple[R, "SparkparseCapture"]:
             with self:
@@ -88,13 +90,14 @@ def __enter__(self):
             ensure_dir(self.temp_dir)
             self._log_dir = self.temp_dir
 
+        orig_conf: dict[str, str] = {}
         if self.spark or SparkSession.getActiveSession():
             self._orig_spark = self.spark
             self._orig_log_dir = self._orig_spark.conf.get("spark.eventLog.dir")
             orig_conf = dict(self._orig_spark.sparkContext._conf.getAll())
             self.spark.stop()
 
-        builder = SparkSession.builder.appName("sparkparse")  # type: ignore
+        builder = SparkSession.builder.appName("sparkparse")
         if hasattr(self, "_orig_spark"):
             for key, value in orig_conf.items():
                 if key not in ["spark.eventLog.enabled", "spark.eventLog.dir"]:
@@ -139,6 +142,10 @@ def _record_history_and_alerts(self) -> None:
         if self._parsed_logs is None or self._history_path is None:
             return
 
+        if self._log_dir is None:
+            _log.warning("log directory is not set, skipping history and alerts")
+            return
+
         effective_log_name = self._log_name or get_path_name(self._log_dir)
         record = history.record_from_dfs(self._parsed_logs, effective_log_name)
         self._last_record = record
@@ -197,7 +204,11 @@ def __exit__(self, exc_type, *args):
         if self._history_path is not None:
             self._record_history_and_alerts()
 
-        if self._should_cleanup and self._log_dir is not None and path_exists(self._log_dir):
+        if (
+            self._should_cleanup
+            and self._log_dir is not None
+            and path_exists(self._log_dir)
+        ):
             remove_dir(self._log_dir)
 
 
@@ -211,7 +222,7 @@ def capture_context(
     alert_config: str | None = None,
 ) -> SparkparseCapture:
     if spark is None:
-        _spark = SparkSession.builder.appName("sparkparse_capture").getOrCreate()  # type: ignore
+        _spark = SparkSession.builder.appName("sparkparse_capture").getOrCreate()
     else:
         _spark = spark
 
@@ -269,7 +280,7 @@ def decorator(
         func: Callable[..., R],
     ) -> Callable[..., tuple[Any, SparkparseCapture]]:
         if spark is None:
-            _spark = SparkSession.builder.appName("sparkparse_capture").getOrCreate()  # type: ignore
+            _spark = SparkSession.builder.appName("sparkparse_capture").getOrCreate()
         else:
             _spark = spark