test: add staged pytest integration suite

Author: marcellodebernardi

AGENTS.md

@@ -81,6 +81,8 @@ make build-databricks         # Databricks Connect
 
 # Run tests
 poetry run pytest tests/unit/
+make test-integration          # Staged pytest integration suite (seed -> search -> eval)
+make test-integration-verbose  # Same suite with live test logs in terminal
 
 # Format and lint
 poetry run black .

CONTRIBUTING.md

@@ -91,6 +91,21 @@ To set up the development environment:
    poetry run pytest
    ```
 
+4. **Run staged integration tests before opening a PR**:
+
+   ```bash
+   # Requires ANTHROPIC_API_KEY and local Spark/Java setup
+   bash scripts/tests/run_integration_staged.sh
+   ```
+
+   The staged suite runs three pytest phases with hard barriers:
+   - `integration_seed`: builds reusable checkpoints through phase 3
+   - `integration_search`: resumes from seeds and runs model search
+   - `integration_eval`: resumes from search checkpoints, runs evaluation, and validates predictor inference
+
+   This `tests/integration` suite is the primary pre-PR integration workflow.
+   Makefile Docker targets remain optional/manual end-to-end checks.
+
 Ensure all tests pass before making contributions.
 
 ## Style Guides

Makefile

@@ -4,6 +4,8 @@
 # Quick reference for developers:
 #   make help              Show all available commands
 #   make test-quick        Fast test (~30s, 1 iteration)
+#   make test-integration  Staged pytest integration suite
+#   make test-integration-verbose  Staged suite with live logs
 #   make test-xgboost      Test XGBoost only
 #   make test-catboost     Test CatBoost only
 #   make test-all-models   Test all model types
@@ -36,6 +38,8 @@ help:
 	@echo "  make test-lightgbm      Test LightGBM model type"
 	@echo "  make test-pytorch       Test PyTorch model type"
 	@echo "  make test-keras         Test Keras model type"
+	@echo "  make test-integration   Run staged pytest integration suite"
+	@echo "  make test-integration-verbose  Run staged suite with live logs"
 	@echo "  make test-all-models    Test all model types (sequential)"
 	@echo "  make test-full          Full test run (3 iterations + evaluation)"
 	@echo ""
@@ -61,6 +65,28 @@ help:
 # Quick Development Tests
 # ============================================
 
+# Staged pytest-native integration suite (seed -> search -> eval).
+# Optional: make test-integration INTEGRATION_RUN_ID=my_run_id
+.PHONY: test-integration
+test-integration:
+	@echo "🧪 Running staged pytest integration suite..."
+	@if [ -n "$(INTEGRATION_RUN_ID)" ]; then \
+		echo "Using integration run id: $(INTEGRATION_RUN_ID)"; \
+		PLEXE_IT_RUN_ID="$(INTEGRATION_RUN_ID)" bash scripts/tests/run_integration_staged.sh; \
+	else \
+		bash scripts/tests/run_integration_staged.sh; \
+	fi
+
+.PHONY: test-integration-verbose
+test-integration-verbose:
+	@echo "🧪 Running staged pytest integration suite (verbose)..."
+	@if [ -n "$(INTEGRATION_RUN_ID)" ]; then \
+		echo "Using integration run id: $(INTEGRATION_RUN_ID)"; \
+		PLEXE_IT_RUN_ID="$(INTEGRATION_RUN_ID)" PLEXE_IT_VERBOSE=1 bash scripts/tests/run_integration_staged.sh; \
+	else \
+		PLEXE_IT_VERBOSE=1 bash scripts/tests/run_integration_staged.sh; \
+	fi
+
 # Fast sanity check - 1 iteration, minimal config
 .PHONY: test-quick
 test-quick: build

poetry.lock

@@ -84,13 +84,21 @@ vision = ["torch"]
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.3.4"
+pytest-xdist = "^3.8.0"
 pre-commit = "^4.0.1"
 ruff = "^0.14.9"
 black = ">=23.0.0"
 streamlit = ">=1.52.1,<2.0.0"
 plotly = ">=6.5.0,<7.0.0"
 boto3 = "^1.42.44"
 
+[tool.pytest.ini_options]
+markers = [
+    "integration_seed: stage 1 integration tests that build reusable checkpoints through phase 3",
+    "integration_search: stage 2 integration tests that resume from seeds and pause after phase 4",
+    "integration_eval: stage 3 integration tests that resume from search checkpoints and run evaluation + packaging",
+]
+
 [tool.semantic_release]
 version_variables = ["pyproject.toml:version"]
 commit_parser = "angular"

pyproject.toml

@@ -0,0 +1,64 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+cd "$ROOT_DIR"
+
+if [[ -z "${PLEXE_IT_RUN_ID:-}" ]]; then
+  PLEXE_IT_RUN_ID="$(date +%Y%m%d_%H%M%S)"
+fi
+export PLEXE_IT_RUN_ID
+
+ARTIFACT_ROOT="$ROOT_DIR/.pytest_cache/integration/$PLEXE_IT_RUN_ID"
+mkdir -p "$ARTIFACT_ROOT"
+
+if ! poetry run python -c "import importlib.util,sys; sys.exit(0 if importlib.util.find_spec('xdist') else 1)"; then
+  echo "ERROR: pytest-xdist is required for staged integration tests."
+  echo "Install dependencies with: poetry install"
+  echo "Then verify with: poetry run pytest --help | grep -E '(^| )-n( |$)'"
+  exit 2
+fi
+
+if [[ -n "${PLEXE_IT_WORKERS:-}" ]]; then
+  WORKERS="${PLEXE_IT_WORKERS}"
+elif [[ "${PLEXE_IT_VERBOSE:-0}" == "1" ]]; then
+  # In verbose mode, default to main-process execution for reliable live logs.
+  WORKERS="0"
+else
+  WORKERS="auto"
+fi
+PYTEST_PARALLEL_ARGS=(-n "$WORKERS")
+
+run_stage() {
+  local stage_name="$1"
+  local marker="$2"
+  local cmd=(poetry run pytest tests/integration -m "$marker" "${PYTEST_PARALLEL_ARGS[@]}" --maxfail=1)
+
+  if [[ "${PLEXE_IT_VERBOSE:-0}" == "1" ]]; then
+    cmd+=(-s -vv -o log_cli=true -o log_cli_level=INFO --capture=tee-sys)
+  fi
+
+  PLEXE_IT_STAGE="$stage_name" "${cmd[@]}"
+}
+
+echo "Running staged integration tests with run id: $PLEXE_IT_RUN_ID"
+echo "Artifacts: $ARTIFACT_ROOT"
+echo "Workers: $WORKERS"
+if [[ "${PLEXE_IT_VERBOSE:-0}" == "1" ]]; then
+  echo "Verbose mode: enabled (live logs and test output)"
+fi
+
+echo ""
+echo "Stage 1/3: building reusable seeds through phase 3"
+run_stage "seed" "integration_seed"
+
+echo ""
+echo "Stage 2/3: resuming from seeds through phase 4"
+run_stage "search" "integration_search"
+
+echo ""
+echo "Stage 3/3: final evaluation, packaging, and predictor checks"
+run_stage "eval" "integration_eval"
+
+echo ""
+echo "Staged integration suite completed successfully."

scripts/tests/run_integration_staged.sh

@@ -4,7 +4,48 @@
 
 Test suite structure and test case documentation.
 
-## `integration/test_feedback.py`
+## `integration/conftest.py`
+Shared fixtures and helpers for staged integration tests.
+
+**Functions:**
+- `repo_root() -> Path` - Return repository root path.
+- `run_id() -> str` - Return deterministic run identifier for staged artifacts.
+- `artifact_root(repo_root: Path, run_id: str) -> Path` - Return base path for staged integration artifacts.
+- `configure_integration_environment(repo_root: Path) -> None` - Set environment variables needed by the integration suite.
+- `cleanup_spark_session() -> None` - Stop Spark session after tests complete.
+- `seed_path(artifact_root: Path, dataset_kind: str) -> Path` - Return seed directory path for a dataset kind.
+- `model_run_path(artifact_root: Path, model_type: str) -> Path` - Return model-specific run directory path.
+- `checkpoint_file(work_dir: Path, phase_name: str) -> Path` - Return path to a checkpoint file.
+- `copy_seed_to_model_run(seed_dir: Path, model_dir: Path) -> None` - Copy a seed workdir into a model run workdir and rewrite checkpoint paths.
+- `assert_stage_prereqs(stage: str, artifact_root: Path) -> None` - Assert required artifacts from prior stages exist.
+- `build_seed_workflow(work_dir: Path, dataset_input: Path, intent: str, experiment_id: str) -> Any` - Run stages 1-3 and pause after baseline creation.
+- `resume_workflow(work_dir: Path, allowed_model_types: list[str], pause_points: list[str] | None, enable_final_evaluation: bool, max_iterations: int) -> Any` - Resume a staged integration workflow from existing checkpoints.
+- `load_predictor_class(model_dir: Path, model_type: str) -> type` - Load predictor class from packaged model/predictor.py.
+- `load_prediction_input(repo_root: Path, dataset_kind: str, n_rows: int) -> pd.DataFrame` - Load a small feature sample used for predictor checks.
+
+---
+## `integration/test_stage1_seed.py`
+Stage 1 integration tests: build reusable checkpoints through phase 3.
+
+**Functions:**
+- `test_build_seed_checkpoint(dataset_kind: str, artifact_root, repo_root) -> None` - Build a seed run and pause after baseline creation.
+
+---
+## `integration/test_stage2_search.py`
+Stage 2 integration tests: resume from seeds and pause after phase 4.
+
+**Functions:**
+- `test_resume_from_seed_and_run_search_only(model_type: str, artifact_root) -> None` - Copy a seed, resume from checkpoints, and pause after search models.
+
+---
+## `integration/test_stage3_eval_predict.py`
+Stage 3 integration tests: run evaluation/packaging and validate predictors.
+
+**Functions:**
+- `test_resume_and_run_eval_then_predict(model_type: str, artifact_root, repo_root) -> None` - Resume from stage 2 checkpoints, run to completion, and validate predictor inference.
+
+---
+## `unit/agents/test_feedback.py`
 Tests for user feedback integration in agents.
 
 **`TestFeedbackFormatting`** - Test the format_user_feedback_for_prompt helper function.