Testing Patterns Skill

Overview

Comprehensive testing patterns for Databricks and Spark applications using pytest, fixtures, and assertion utilities. This skill provides reusable patterns for unit testing transformations, integration testing pipelines, and mocking Spark components.

Key Benefits:

• Reusable test fixtures for Spark

• DataFrame assertion utilities

• Mock data generation patterns

• Integration test strategies

• CI/CD test automation

• Coverage tracking

When to Use This Skill

Use testing patterns when you need to:

• Unit test data transformations

• Validate pipeline logic before deployment

• Set up regression testing

• Create test data for development

• Implement CI/CD quality gates

• Debug transformation logic

• Document expected behavior

Core Concepts

1. Pytest Configuration

conftest.py Setup:

import pytest from pyspark.sql import SparkSession

@pytest.fixture(scope="session") def spark(): """Create shared Spark session for all tests.""" return ( SparkSession.builder .master("local[2]") .appName("pytest-spark") .config("spark.sql.shuffle.partitions", "2") .getOrCreate() )

@pytest.fixture def sample_customers(spark): """Generate sample customer data.""" return spark.createDataFrame([ ("C1", "Alice", "alice@example.com", 100.0), ("C2", "Bob", "bob@example.com", 200.0), ("C3", "Charlie", "charlie@example.com", 300.0) ], ["customer_id", "name", "email", "lifetime_value"])

2. DataFrame Assertions

Custom Assertion Helpers:

from pyspark.sql import DataFrame from typing import List, Set

def assert_schema_match(df: DataFrame, expected_columns: List[str]): """Assert DataFrame has expected schema.""" actual_columns = df.columns assert set(actual_columns) == set(expected_columns), \ f"Schema mismatch: got {actual_columns}, expected {expected_columns}"

def assert_no_nulls(df: DataFrame, columns: List[str]): """Assert specified columns have no null values.""" for col in columns: null_count = df.filter(df[col].isNull()).count() assert null_count == 0, f"Column '{col}' has {null_count} null values"

def assert_unique(df: DataFrame, key_columns: List[str]): """Assert key columns form unique keys.""" total_count = df.count() unique_count = df.select(key_columns).distinct().count() assert total_count == unique_count, \ f"Duplicate keys found: {total_count - unique_count} duplicates"

3. Mock Data Generation

Test Data Builders:

from pyspark.sql import SparkSession from pyspark.sql.types import * from datetime import datetime, timedelta import random

class TestDataBuilder: """Build realistic test data."""

def __init__(self, spark: SparkSession):
    self.spark = spark

def generate_orders(self, count: int = 100):
    """Generate sample order data."""
    schema = StructType([
        StructField("order_id", StringType(), False),
        StructField("customer_id", StringType(), False),
        StructField("order_date", TimestampType(), False),
        StructField("amount", DoubleType(), False),
        StructField("status", StringType(), False)
    ])

    data = [
        (
            f"O{i:05d}",
            f"C{random.randint(1, 50):03d}",
            datetime.now() - timedelta(days=random.randint(0, 365)),
            round(random.uniform(10.0, 1000.0), 2),
            random.choice(["pending", "completed", "cancelled"])
        )
        for i in range(count)
    ]

    return self.spark.createDataFrame(data, schema)

4. Integration Testing

End-to-End Pipeline Tests:

import pytest from databricks.sdk import WorkspaceClient from databricks.sdk.service.jobs import RunResultState

@pytest.mark.integration class TestPipelineIntegration: """Integration tests for complete pipeline."""

@pytest.fixture
def workspace_client(self):
    """Create Databricks workspace client."""
    return WorkspaceClient()

def test_bronze_to_silver_pipeline(
    self,
    spark,
    workspace_client,
    tmp_path
):
    """Test bronze to silver transformation."""
    # Arrange: Create test data
    bronze_data = spark.createDataFrame([
        ("1", "test@example.com", "100"),
        ("2", "invalid_email", "200"),
        ("3", None, "300")
    ], ["id", "email", "amount"])

    bronze_path = str(tmp_path / "bronze")
    bronze_data.write.format("delta").save(bronze_path)

    # Act: Run transformation
    silver_df = transform_bronze_to_silver(
        spark.read.format("delta").load(bronze_path)
    )

    # Assert: Verify results
    assert silver_df.count() == 2  # Invalid record dropped
    assert_no_nulls(silver_df, ["id", "email"])
    assert_schema_match(
        silver_df,
        ["id", "email", "amount", "validated_at"]
    )

Implementation Patterns

Pattern 1: Unit Testing Transformations

Test Data Transformations:

""" Unit tests for order transformations. """ import pytest from pyspark.sql.functions import col, when, lit from datetime import datetime

def enrich_orders(df): """Add calculated fields to orders.""" return df.withColumn( "discount_amount", when(col("amount") > 100, col("amount") * 0.1).otherwise(lit(0.0)) ).withColumn( "enriched_at", lit(datetime.now()) )

class TestOrderTransformations: """Test order transformation logic."""

def test_enrich_orders_adds_discount(self, spark):
    """Test discount calculation logic."""
    # Arrange
    input_df = spark.createDataFrame([
        ("O1", 150.0),
        ("O2", 50.0)
    ], ["order_id", "amount"])

    # Act
    result = enrich_orders(input_df)

    # Assert
    assert "discount_amount" in result.columns
    high_order = result.filter(col("order_id") == "O1").first()
    low_order = result.filter(col("order_id") == "O2").first()

    assert high_order["discount_amount"] == 15.0
    assert low_order["discount_amount"] == 0.0

def test_enrich_orders_adds_timestamp(self, spark):
    """Test timestamp addition."""
    input_df = spark.createDataFrame([
        ("O1", 100.0)
    ], ["order_id", "amount"])

    result = enrich_orders(input_df)

    assert "enriched_at" in result.columns
    assert result.first()["enriched_at"] is not None

def test_enrich_orders_preserves_original_columns(self, spark):
    """Test original columns are preserved."""
    input_df = spark.createDataFrame([
        ("O1", 100.0)
    ], ["order_id", "amount"])

    result = enrich_orders(input_df)

    assert set(input_df.columns).issubset(set(result.columns))

Pattern 2: Parameterized Testing

Test Multiple Scenarios:

import pytest

@pytest.mark.parametrize("amount,expected_discount", [ (50.0, 0.0), # Below threshold (100.0, 0.0), # At threshold (150.0, 15.0), # Above threshold (1000.0, 100.0) # High value ]) def test_discount_calculation(spark, amount, expected_discount): """Test discount calculation with various amounts.""" df = spark.createDataFrame([ ("O1", amount) ], ["order_id", "amount"])

result = enrich_orders(df)

actual_discount = result.first()["discount_amount"]
assert actual_discount == expected_discount

@pytest.mark.parametrize("email,expected_valid", [ ("user@example.com", True), ("invalid_email", False), ("", False), (None, False) ]) def test_email_validation(spark, email, expected_valid): """Test email validation logic.""" df = spark.createDataFrame([ ("C1", email) ], ["customer_id", "email"])

result = validate_customer_data(df)

is_valid = result.filter(col("customer_id") == "C1").count() > 0
assert is_valid == expected_valid

Pattern 3: Fixture Factories

Reusable Test Data Factories:

""" Test data factories for common entities. """ import pytest from datetime import datetime, timedelta from typing import List, Dict, Any

@pytest.fixture def customer_factory(spark): """Factory for generating customer test data.""" def _create_customers( count: int = 10, status: str = "active", **overrides ): base_data = [ { "customer_id": f"C{i:05d}", "name": f"Customer {i}", "email": f"customer{i}@example.com", "status": status, "created_at": datetime.now() - timedelta(days=i), **overrides } for i in range(count) ] return spark.createDataFrame(base_data)

return _create_customers

@pytest.fixture def order_factory(spark): """Factory for generating order test data.""" def _create_orders( customer_ids: List[str], orders_per_customer: int = 5 ): data = [] order_num = 0 for customer_id in customer_ids: for i in range(orders_per_customer): data.append({ "order_id": f"O{order_num:06d}", "customer_id": customer_id, "amount": 100.0 * (i + 1), "order_date": datetime.now() - timedelta(days=i) }) order_num += 1 return spark.createDataFrame(data)

return _create_orders

Usage in tests

def test_customer_aggregation(customer_factory, order_factory): """Test customer order aggregation.""" customers = customer_factory(count=3) customer_ids = [row["customer_id"] for row in customers.collect()] orders = order_factory(customer_ids, orders_per_customer=5)

result = aggregate_customer_orders(customers, orders)

assert result.count() == 3
assert "total_orders" in result.columns
assert "total_amount" in result.columns

Pattern 4: DLT Pipeline Testing

Test DLT Functions:

""" Test DLT pipeline functions. """ import pytest from unittest.mock import Mock, patch

def test_dlt_bronze_ingestion(spark, tmp_path): """Test bronze layer ingestion logic.""" # Mock DLT read_stream with patch('dlt.read_stream') as mock_read: # Create test data test_data = spark.createDataFrame([ ("1", "data1"), ("2", "data2") ], ["id", "value"])

    mock_read.return_value = test_data

    # Import and test the transformation function
    from pipelines.bronze import transform_raw_data

    result = transform_raw_data()

    assert "ingested_at" in result.columns
    assert result.count() == 2

def test_dlt_expectations_logic(spark): """Test data quality logic without DLT decorators.""" # Extract core transformation logic def apply_quality_rules(df): """Core quality logic from DLT pipeline.""" return df.filter( (col("customer_id").isNotNull()) & (col("amount") > 0) )

# Test data with violations
test_data = spark.createDataFrame([
    ("C1", 100.0),   # Valid
    (None, 200.0),   # Null customer_id
    ("C2", -50.0),   # Negative amount
    ("C3", 150.0)    # Valid
], ["customer_id", "amount"])

result = apply_quality_rules(test_data)

assert result.count() == 2  # Only valid records

Pattern 5: Performance Testing

Test Performance Characteristics:

""" Performance tests for transformations. """ import pytest import time

@pytest.mark.performance def test_large_dataset_performance(spark): """Test transformation performance with large dataset.""" # Generate large test dataset large_df = spark.range(0, 1000000).toDF("id") large_df = large_df.withColumn("value", col("id") * 2)

# Measure transformation time
start_time = time.time()

result = complex_transformation(large_df)
result.write.format("noop").mode("overwrite").save()

duration = time.time() - start_time

# Assert performance requirements
assert duration < 60, f"Transformation took {duration}s, expected < 60s"
assert result.count() == 1000000

@pytest.mark.performance def test_join_performance(spark, customer_factory, order_factory): """Test join performance with realistic data volumes.""" customers = customer_factory(count=10000) customer_ids = [row["customer_id"] for row in customers.limit(1000).collect()] orders = order_factory(customer_ids, orders_per_customer=100)

start_time = time.time()

result = customers.join(orders, "customer_id", "inner")
result.write.format("noop").mode("overwrite").save()

duration = time.time() - start_time

assert duration < 30, f"Join took {duration}s, expected < 30s"

Best Practices

1. Test Organization

tests/ ├── unit/ │ ├── conftest.py │ ├── test_transformations.py │ ├── test_validators.py │ └── test_aggregations.py ├── integration/ │ ├── conftest.py │ ├── test_pipeline_end_to_end.py │ └── test_dlt_integration.py ├── fixtures/ │ ├── sample_data.py │ └── mock_services.py └── performance/ └── test_performance.py

2. Naming Conventions

• Test files: test_*.py

• Test functions: test_<what_is_tested>

• Fixtures: descriptive nouns (e.g., customer_data , spark_session )

• Test classes: Test<Component> (e.g., TestOrderTransformations )

3. Test Independence

Good: Each test is independent

def test_transformation_a(spark): df = spark.createDataFrame([...]) result = transform_a(df) assert result.count() == expected

def test_transformation_b(spark): df = spark.createDataFrame([...]) # Fresh data result = transform_b(df) assert result.count() == expected

Bad: Tests depend on shared state

shared_df = None

def test_transformation_a(spark): global shared_df shared_df = transform_a(spark.createDataFrame([...]))

def test_transformation_b(): result = transform_b(shared_df) # Depends on test_a

4. Assertion Quality

Good: Specific assertions

def test_discount_calculation(spark): result = calculate_discount(order_df) assert result.filter(col("amount") > 100).select("discount").first()[0] == 10.0

Bad: Weak assertions

def test_discount_calculation(spark): result = calculate_discount(order_df) assert result is not None # Too generic assert result.count() > 0 # Doesn't test logic

5. Mock External Dependencies

from unittest.mock import patch, Mock

def test_external_api_call(spark): """Test transformation that calls external API.""" with patch('requests.get') as mock_get: mock_response = Mock() mock_response.json.return_value = {"rate": 1.2} mock_get.return_value = mock_response

    result = enrich_with_exchange_rate(order_df)

    mock_get.assert_called_once()
    assert "exchange_rate" in result.columns

Common Pitfalls to Avoid

Don't:

• Test framework code (e.g., testing that Spark works)

• Write tests that depend on execution order

• Use production data in tests

• Skip cleanup of test resources

• Mix unit and integration tests

Do:

• Test business logic and transformations

• Make tests independent and idempotent

• Generate synthetic test data

• Clean up temporary resources

• Separate test types (unit, integration, performance)

Complete Examples

See /examples/ directory for:

• test_transformation_pipeline.py : Complete unit test suite

• test_integration_workflow.py : End-to-end integration tests

Related Skills

• delta-live-tables : Testing DLT pipelines

• data-quality : Quality validation testing

• cicd-workflows : Automated test execution

• medallion-architecture : Layer-specific testing

References

• Pytest Documentation

• PySpark Testing Best Practices

• Databricks Testing Guide