PySpark 코드 테스트 — 데이터 파이프라인을 회귀로부터 지키기

# BAD: I/O 와 로직이 뒤섞여 테스트 불가
def job():
    df = spark.read.parquet("s3://...")
    result = df.filter(...).groupBy(...).agg(...)   # 로직이 여기 묶임
    result.write.parquet("s3://...")
 
# GOOD: 변환을 순수 함수로 분리
def transform(df):                    # 입력 DataFrame → 출력 DataFrame
    return (df.filter(F.col("amount") > 0)
              .groupBy("user").agg(F.sum("amount").alias("total")))
 
def job():
    df = spark.read.parquet("s3://...")
    transform(df).write.parquet("s3://...")   # I/O 는 얇게

# conftest.py
import pytest
from pyspark.sql import SparkSession
 
@pytest.fixture(scope="session")
def spark():
    spark = (SparkSession.builder
        .master("local[2]")                         # 로컬 2코어
        .appName("tests")
        .config("spark.sql.shuffle.partitions", "2") # 테스트는 작게 → 빠르게
        .config("spark.ui.enabled", "false")
        .getOrCreate())
    yield spark
    spark.stop()

from pyspark.testing import assertDataFrameEqual
 
def test_transform(spark):
    input_df = spark.createDataFrame(
        [("a", 100), ("a", 50), ("b", 30), ("b", -10)],  # -10 은 필터돼야
        ["user", "amount"])
 
    expected = spark.createDataFrame(
        [("a", 150), ("b", 30)], ["user", "total"])
 
    result = transform(input_df)
    assertDataFrameEqual(result, expected)   # 순서 무관 비교

from chispa import assert_df_equality
 
assert_df_equality(result, expected, ignore_row_order=True, ignore_nullable=True)

def test_empty_input(spark):
    # 빈 DataFrame 도 깨지지 않아야
    empty = spark.createDataFrame([], "user string, amount long")
    result = transform(empty)
    assert result.count() == 0
 
def test_nulls(spark):
    # NULL 처리 — 집계·필터에서 NULL 이 어떻게 되는가
    df = spark.createDataFrame([("a", None), ("a", 100)], "user string, amount long")
    result = transform(df)
    # 기대 동작을 명시적으로 검증
 
def test_negative_and_zero(spark):
    # 경계값: 0, 음수
    ...

def test_output_schema(spark):
    input_df = spark.createDataFrame([("a", 100)], ["user", "amount"])
    result = transform(input_df)
 
    expected_schema = ["user", "total"]
    assert result.columns == expected_schema
    assert dict(result.dtypes)["total"] == "bigint"

def assert_quality(df):
    total = df.count()
    # 핵심 컬럼 NULL 없어야
    null_ids = df.filter(F.col("user").isNull()).count()
    assert null_ids == 0, f"user NULL {null_ids}건"
    # 키 유일성
    assert df.select("user").distinct().count() == total, "user 중복"
    # 값 범위
    assert df.filter(F.col("total") < 0).count() == 0, "음수 total"

# 예: GitHub Actions 개념
steps:
  - uses: actions/setup-java@v4   # Spark 는 JVM 필요
    with: { java-version: '17' }
  - run: pip install pyspark chispa pytest
  - run: pytest tests/ -v

        ▲  통합 테스트 (실제 읽기/쓰기, 소수)
       ───
      ─────  변환 단위 테스트 (transform 함수, 다수)  ← 여기에 집중
     ───────
    ─────────  순수 헬퍼 함수 (Spark 불필요, 가장 빠름)