12.4 Automated Testing, Unit Testing PySpark, & CI/CD Pipelines

12.4 Automated Testing, Unit Testing PySpark, & CI/CD Pipelines

A resilient data platform requires automated testing and continuous integration to catch data processing bugs, schema regressions, and syntax errors before code reaches production lakehouses. Deploying untested transformation notebooks directly to production often leads to silent data corruption, downstream dashboard failures, and expensive cluster reprocessing.

This section covers the end-to-end automation lifecycle: structuring testable PySpark code, executing local unit tests using pytest and chispa, mocking input datasets, and configuring production-grade GitHub Actions and Azure DevOps Pipelines using Databricks Asset Bundles and Microsoft Entra Service Principal authentication.

---

1. Structuring Testable PySpark Code

The primary barrier to unit testing data engineering pipelines is the coupling of business transformation logic with data ingestion and egress (reading from and writing to tables/files).

+-------------------------------------------------------------------------+
|                   POOR DESIGN: COUPLED I/O & LOGIC                      |
+-------------------------------------------------------------------------+
|  def process_sales():                                                   |
|      df = spark.read.table("prod.bronze.sales")  # Hardcoded I/O        |
|      df_clean = df.filter(df.amount > 0)         # Transformation       |
|      df_clean.write.saveAsTable("prod.silver.sales") # Hardcoded I/O    |
|  * Impossible to unit test without connecting to live production data!  |
+-------------------------------------------------------------------------+
|                 RECOMMENDED DESIGN: DECOUPLED PURE FUNCTIONS            |
+-------------------------------------------------------------------------+
|  # Pure transformation function: Accepts DataFrame -> Returns DataFrame |
|  def transform_sales_data(input_df: DataFrame) -> DataFrame:            |
|      return input_df.filter(input_df.amount > 0)                        |
|  * Can be tested instantly in local pytest using mock DataFrames!       |
+-------------------------------------------------------------------------+

Implementation: Modular Transformation Library (src/transforms.py)

# File: src/transforms.py
from pyspark.sql import DataFrame
from pyspark.sql import functions as F

def calculate_customer_metrics(transactions_df: DataFrame) -> DataFrame:
    """
    Cleanses transactions, calculates net amount, and classifies tier.
    """
    return (
        transactions_df
        .filter(F.col("transaction_id").isNotNull())
        .filter(F.col("amount") > 0)
        .withColumn("tax", F.round(F.col("amount") * 0.08, 2))
        .withColumn("net_total", F.col("amount") + F.col("tax"))
        .withColumn(
            "customer_tier",
            F.when(F.col("net_total") >= 1000, "VIP")
            .when(F.col("net_total") >= 250, "PREMIUM")
            .otherwise("STANDARD")
        )
    )

---

2. Unit Testing PySpark with pytest and chispa

Testing PySpark DataFrames requires comparing distributed datasets, schema definitions, nullable flags, and floating-point numeric precision.

Setting up the Test Environment (tests/conftest.py)

A shared pytest fixture creates an in-memory, local SparkSession on the test runner without connecting to cloud clusters:

# File: tests/conftest.py
import pytest
from pyspark.sql import SparkSession

@pytest.fixture(scope="session")
def spark():
    """Provides a local SparkSession for unit tests."""
    spark_session = (
        SparkSession.builder
        .master("local[2]")
        .appName("PySparkUnitTests")
        .config("spark.sql.shuffle.partitions", "2")
        .config("spark.default.parallelism", "2")
        .getOrCreate()
    )
    yield spark_session
    spark_session.stop()

Writing Tests with chispa and assertDataFrameEqual (tests/test_transforms.py)

chispa is an industry-standard PySpark testing library that outputs detailed, colorized column and row diffs upon test failure:

# File: tests/test_transforms.py
import pytest
from pyspark.sql.types import StructType, StructField, StringType, DoubleType
from chispa.dataframe_comparer import assert_df_equality
from src.transforms import calculate_customer_metrics

def test_calculate_customer_metrics(spark):
    # 1. Arrange: Define mock input schema and test data
    input_schema = StructType([
        StructField("transaction_id", StringType(), True),
        StructField("customer_id", StringType(), True),
        StructField("amount", DoubleType(), True)
    ])
    
    input_data = [
        ("tx_001", "cust_101", 100.0),    # Standard tier (net: 108.0)
        ("tx_002", "cust_102", 500.0),    # Premium tier  (net: 540.0)
        ("tx_003", "cust_103", 2000.0),   # VIP tier      (net: 2160.0)
        ("tx_004", "cust_104", -50.0),    # Negative amount (should be filtered out)
        (None, "cust_105", 300.0)         # Null transaction_id (filtered out)
    ]
    input_df = spark.createDataFrame(input_data, input_schema)
    
    # 2. Act: Execute transformation function
    result_df = calculate_customer_metrics(input_df)
    
    # 3. Assert: Define expected output dataset
    expected_schema = StructType([
        StructField("transaction_id", StringType(), True),
        StructField("customer_id", StringType(), True),
        StructField("amount", DoubleType(), True),
        StructField("tax", DoubleType(), True),
        StructField("net_total", DoubleType(), True),
        StructField("customer_tier", StringType(), False)
    ])
    
    expected_data = [
        ("tx_001", "cust_101", 100.0, 8.0, 108.0, "STANDARD"),
        ("tx_002", "cust_102", 500.0, 40.0, 540.0, "PREMIUM"),
        ("tx_003", "cust_103", 2000.0, 160.0, 2160.0, "VIP")
    ]
    expected_df = spark.createDataFrame(expected_data, expected_schema)
    
    # Compare DataFrames ignoring row order and minor float rounding differences
    assert_df_equality(
        result_df,
        expected_df,
        ignore_row_order=True,
        ignore_nullable=True
    )

Exam Tip: In PySpark 3.5+ (Databricks Runtime 14+), PySpark includes the built-in testing function pyspark.testing.assertDataFrameEqual(actual, expected, checkRowOrder=False, rtol=1e-5) which can be used without external third-party dependencies.

---

3. End-to-End CI/CD Pipeline Architecture

A mature Lakehouse CI/CD pipeline implements a multi-stage promotion gate:

+-------------------------------------------------------------------------+
|                    LAKEHOUSE CI/CD PROMOTION LIFECYCLE                  |
+-------------------------------------------------------------------------+
|                                                                         |
|  1. PULL REQUEST (PR GATE)                                              |
|     ├── Code Linting (flake8, ruff, black)                              |
|     ├── Local PySpark Unit Tests (pytest + chispa on GitHub runner)     |
|     └── DAB Validation (databricks bundle validate -t staging)          |
|                                                                         |
|  2. MERGE TO MAIN (STAGING PROMOTION)                                   |
|     ├── Deploy Bundle to Staging (databricks bundle deploy -t staging)  |
|     └── Run Integration Test Pipeline (databricks bundle run -t staging)|
|                                                                         |
|  3. RELEASE TAG / APPROVAL (PRODUCTION DEPLOYMENT)                      |
|     ├── Environment Approval Gate (Manual or Automated Sign-off)        |
|     └── Deploy to Prod via Service Principal (bundle deploy -t prod)   |
+-------------------------------------------------------------------------+

---

4. GitHub Actions CI/CD Pipeline Implementation

The following workflow file (.github/workflows/databricks_ci_cd.yml) automates testing and deployment:

name: Lakehouse CI/CD Pipeline

on:
  pull_request:
    branches: [main]
  push:
    branches: [main]

env:
  DATABRICKS_HOST: ${{ secrets.DATABRICKS_HOST }}
  DATABRICKS_CLIENT_ID: ${{ secrets.AZURE_SP_CLIENT_ID }}
  DATABRICKS_CLIENT_SECRET: ${{ secrets.AZURE_SP_CLIENT_SECRET }}

jobs:
  # STAGE 1: Fast Unit Testing & Linting
  unit_tests:
    name: Run PySpark Unit Tests
    runs-on: ubuntu-latest
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4

      - name: Set up Python 3.11
        uses: actions/setup-python@v5
        with:
          python-version: '3.11'
          cache: 'pip'

      - name: Set up Java JDK 17 (Required for local Spark)
        uses: actions/setup-java@v4
        with:
          distribution: 'temurin'
          java-version: '17'

      - name: Install Dependencies
        run: |
          python -m pip install --upgrade pip
          pip install pytest chispa pyspark==3.5.1 ruff

      - name: Run Ruff Linting
        run: ruff check src/ tests/

      - name: Run PyTest Suite
        run: pytest tests/ -v

  # STAGE 2: DAB Validation & Staging Deployment
  deploy_staging:
    name: Validate & Deploy to Staging
    needs: unit_tests
    runs-on: ubuntu-latest
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4

      - name: Setup Databricks CLI
        uses: databricks/setup-cli@v0.2

      - name: Validate Bundle for Staging
        run: databricks bundle validate -t staging

      - name: Deploy to Staging Workspace
        run: databricks bundle deploy -t staging

      - name: Run Integration Smoke Test in Staging
        run: databricks bundle run -t staging sales_daily_etl_job

  # STAGE 3: Production Deployment (On Merge to Main only)
  deploy_production:
    name: Deploy to Production
    needs: deploy_staging
    if: github.ref == 'refs/heads/main' && github.event_name == 'push'
    runs-on: ubuntu-latest
    environment: production
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4

      - name: Setup Databricks CLI
        uses: databricks/setup-cli@v0.2

      - name: Validate Bundle for Production
        run: databricks bundle validate -t prod

      - name: Deploy to Production Workspace
        run: databricks bundle deploy -t prod

---

5. Azure DevOps Release Pipeline Implementation

For enterprise Azure environments, the pipeline is configured in azure-pipelines.yml:

trigger:
  branches:
    include:
      - main

pool:
  vmImage: 'ubuntu-latest'

variables:
  - group: databricks-prod-secrets # Azure Key Vault backed variable group

stages:
  - stage: TestAndValidate
    displayName: 'Unit Tests & DAB Validate'
    jobs:
      - job: PyTest
        steps:
          - task: UsePythonVersion@0
            inputs:
              versionSpec: '3.11'
          - task: JavaToolInstaller@0
            inputs:
              versionSpec: '17'
              jdkArchitectureOption: 'x64'
              jdkSourceOption: 'PreInstalled'
          - script: |
              pip install pytest chispa pyspark
              pytest tests/
            displayName: 'Execute PySpark Unit Tests'
          - script: |
              curl -fsSL https://raw.githubusercontent.com/databricks/setup-cli/main/install.sh | sh
              databricks bundle validate -t prod
            displayName: 'Validate DAB Manifest'
            env:
              DATABRICKS_HOST: $(DATABRICKS_HOST)
              DATABRICKS_CLIENT_ID: $(AZURE_SP_CLIENT_ID)
              DATABRICKS_CLIENT_SECRET: $(AZURE_SP_CLIENT_SECRET)

  - stage: DeployProd
    displayName: 'Deploy to Production Workspace'
    dependsOn: TestAndValidate
    condition: and(succeeded(), eq(variables['Build.SourceBranch'], 'refs/heads/main'))
    jobs:
      - deployment: DeployDAB
        environment: 'Databricks-Production'
        strategy:
          runOnce:
            deploy:
              steps:
                - checkout: self
                - script: |
                    curl -fsSL https://raw.githubusercontent.com/databricks/setup-cli/main/install.sh | sh
                    databricks bundle deploy -t prod
                  displayName: 'Deploy DAB to Production Workspace'
                  env:
                    DATABRICKS_HOST: $(DATABRICKS_HOST)
                    DATABRICKS_CLIENT_ID: $(AZURE_SP_CLIENT_ID)
                    DATABRICKS_CLIENT_SECRET: $(AZURE_SP_CLIENT_SECRET)