Тестирование PySpark: от хаоса к надёжности

# Монолитный пайплайн - невозможно тестировать изолированно
def run_pipeline():
    spark = SparkSession.builder.getOrCreate()

    # Чтение из S3 - без него пайплайн не запустится даже ради одного теста
    df = spark.read.parquet("s3a://bucket/raw/transactions/")

    # Трансформации перемешаны с зависимостями на внешние системы
    df = df.filter(col("status") == "completed")
    df = df.withColumn("tax", col("amount") * 0.2)
    df = df.join(
        spark.read.parquet("s3a://bucket/dim/users/"),
        "user_id"
    )

    df.write.format("delta").mode("overwrite").save("s3a://bucket/silver/")

# transformations.py - только чистые функции, никакого IO
from pyspark.sql import DataFrame
from pyspark.sql import functions as F


def filter_completed_transactions(df: DataFrame) -> DataFrame:
    """Оставляет только транзакции со статусом 'completed'."""
    return df.filter(F.col("status") == "completed")


def calculate_tax(df: DataFrame, tax_rate: float = 0.2) -> DataFrame:
    """
    Рассчитывает налог как долю от суммы.
    coalesce защищает от NULL: NULL * 0.2 = NULL, но нам нужен 0.0.
    """
    return df.withColumn(
        "tax",
        F.coalesce(F.col("amount"), F.lit(0.0)) * tax_rate
    )


def add_discount_flag(df: DataFrame) -> DataFrame:
    """Помечает транзакции, где скидка превышает 50% от суммы."""
    return df.withColumn(
        "big_discount",
        F.when(
            F.col("discount").isNull() | (F.col("amount") <= 0),
            F.lit(False)
        ).otherwise(
            F.col("discount") > (F.col("amount") * 0.5)
        )
    )

# pipeline.py - оркестрация: только IO + вызовы трансформаций
from pyspark.sql import SparkSession
from transformations import (
    filter_completed_transactions,
    calculate_tax,
    add_discount_flag,
)


def run(spark: SparkSession, input_path: str, output_path: str) -> None:
    df = spark.read.parquet(input_path)

    result = (
        df
        .transform(filter_completed_transactions)
        .transform(calculate_tax)
        .transform(add_discount_flag)
    )

    result.write.format("delta").mode("overwrite").save(output_path)

pip install pytest chispa pyspark pytest-cov

# Для работы Spark нужна Java 8 или новее
java -version
# java version "17.0.8" 2023-07-18 - это нормально

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


@pytest.fixture(scope="session")
def spark():
    """SparkSession в local-режиме для unit-тестов."""
    session = (
        SparkSession.builder
        .master("local[*]")                               # все доступные ядра CPU
        .appName("pyspark-unit-tests")
        .config("spark.sql.shuffle.partitions", "4")      # мало партиций = быстрее
        .config("spark.ui.enabled", "false")              # не открывать web UI
        .config("spark.sql.adaptive.enabled", "false")    # AQE не нужен в тестах
        .config("spark.driver.log.level", "WARN")         # меньше логов в консоль
        .getOrCreate()
    )
    yield session
    session.stop()

my_pipeline/
├── src/
│   ├── transformations.py       # чистые трансформации: DataFrame → DataFrame
│   ├── udfs.py                  # пользовательские UDF
│   └── pipeline.py              # оркестрация: IO + вызовы трансформаций
├── tests/
│   ├── conftest.py              # SparkSession fixture, общие данные
│   ├── test_transformations.py  # unit-тесты трансформаций
│   ├── test_udfs.py             # тесты UDF
│   ├── test_aggregations.py     # тесты агрегаций
│   ├── test_joins.py            # тесты join-логики
│   └── test_edge_cases.py       # NULL, пустые DF, граничные значения
├── .github/workflows/test.yml   # CI/CD
└── pyproject.toml               # зависимости и настройки pytest

pytest tests/ -v                                          # все тесты
pytest tests/test_transformations.py -v                   # один файл
pytest tests/ -k "tax" -v                                 # тесты с "tax" в имени
pytest tests/test_transformations.py::test_calculate_tax  # конкретный тест
pytest tests/ -v --cov=src --cov-report=term-missing      # с отчётом coverage

# Сравнивает Python-объекты, а не содержимое - всегда False или ошибка
assert result == expected

# Чувствительно к порядку строк - упадёт при перемешивании партиций
assert result.collect() == expected.collect()

# Не показывает ЧТО именно отличается - знаем только что "не равно"
assert result.collect() == expected.collect()
# AssertionError. Где ошибка? Не понятно.

from chispa.dataframe_comparer import assert_df_equality
from src.transformations import calculate_tax


def test_calculate_tax_standard_rate(spark):
    """Расчёт налога 20% для обычных транзакций."""
    input_df = spark.createDataFrame(
        [(1, 100.0), (2, 250.0), (3, 50.0)],
        ["transaction_id", "amount"]
    )

    result = calculate_tax(input_df, tax_rate=0.2)

    expected = spark.createDataFrame(
        [(1, 100.0, 20.0), (2, 250.0, 50.0), (3, 50.0, 10.0)],
        ["transaction_id", "amount", "tax"]
    )

    # ignore_row_order=True - обязательный флаг для большинства тестов
    assert_df_equality(result, expected, ignore_row_order=True)

DataFramesNotEqualError:
Left DataFrame:
+---+------+----+
| id|amount| tax|
+---+------+----+
|  1| 100.0|20.0|
|  2| 250.0|52.0|  ← Ожидалось 50.0, но получили 52.0
+---+------+----+

Right DataFrame:
+---+------+----+
| id|amount| tax|
+---+------+----+
|  1| 100.0|20.0|
|  2| 250.0|50.0|
+---+------+----+

from chispa.column_comparer import assert_column_equality


def test_tax_column_values(spark):
    """Проверяем только значения колонки tax, игнорируя остальные."""
    input_df = spark.createDataFrame(
        [(1, 100.0), (2, 200.0)],
        ["id", "amount"]
    )
    result = calculate_tax(input_df)

    expected_series = spark.createDataFrame(
        [(20.0,), (40.0,)],
        ["tax"]
    )
    # Сравниваем только колонку "tax" из result с колонкой "tax" из expected_series
    assert_column_equality(result, expected_series, "tax")

from pyspark.testing import assertDataFrameEqual


def test_with_builtin(spark):
    result = calculate_tax(input_df)
    assertDataFrameEqual(result, expected, checkRowOrder=False)

# Самый короткий способ: список кортежей + список имён
df = spark.createDataFrame(
    [(1, "completed", 100.0), (2, "pending", 50.0)],
    ["id", "status", "amount"]
)

from pyspark.sql.types import (
    StructType, StructField,
    IntegerType, StringType, DoubleType,
)


def test_with_explicit_schema(spark):
    schema = StructType([
        StructField("transaction_id", IntegerType(), nullable=False),
        StructField("amount",         DoubleType(),  nullable=True),
        StructField("status",         StringType(),  nullable=True),
    ])

    input_df = spark.createDataFrame(
        [(1, 100.0, "completed"), (2, None, "completed")],
        schema
    )

    result = calculate_tax(input_df)

    # Убеждаемся, что tax имеет нужный тип
    assert result.schema["tax"].dataType == DoubleType()

# conftest.py
@pytest.fixture
def transactions_df(spark):
    """Базовый DataFrame транзакций: типичные, NULL и граничные случаи."""
    return spark.createDataFrame(
        [
            (1, "completed", 100.0, 10.0),
            (2, "pending",   250.0,  0.0),
            (3, "completed",  None,   0.0),  # NULL amount
            (4, "completed",   0.0,   5.0),  # нулевая сумма со скидкой
        ],
        ["id", "status", "amount", "discount"]
    )

import pytest
from chispa.dataframe_comparer import assert_df_equality
from src.transformations import calculate_tax


@pytest.mark.parametrize("tax_rate, amount, expected_tax", [
    (0.20, 100.0, 20.0),    # стандартная ставка
    (0.10, 100.0, 10.0),    # льготная ставка
    (0.0,  100.0,  0.0),    # нулевая ставка (освобождение от налога)
    (0.20, 250.0, 50.0),    # большая сумма
    (0.20,   0.0,  0.0),    # нулевая сумма
])
def test_calculate_tax_rates(spark, tax_rate, amount, expected_tax):
    """Параметризованный тест: проверяем разные ставки и суммы за один прогон."""
    df = spark.createDataFrame(
        [(1, amount)],
        ["id", "amount"]
    )
    result = calculate_tax(df, tax_rate=tax_rate)
    actual = result.collect()[0]["tax"]

    # Для float всегда используем допустимую погрешность
    assert abs(actual - expected_tax) < 0.001, \
        f"tax_rate={tax_rate}, amount={amount}: ожидалось {expected_tax}, получено {actual}"

test_calculate_tax_rates[0.2-100.0-20.0] PASSED
test_calculate_tax_rates[0.1-100.0-10.0] PASSED
test_calculate_tax_rates[0.0-100.0-0.0]  PASSED
test_calculate_tax_rates[0.2-250.0-50.0] PASSED
test_calculate_tax_rates[0.2-0.0-0.0]    PASSED

@pytest.mark.parametrize("status, expected_count", [
    ("completed",  2),   # стандартный статус
    ("COMPLETED",  0),   # регистр важен - этот не пройдёт фильтр
    ("cancelled",  0),   # отменённые не включаются
    ("pending",    0),   # ожидающие не включаются
    (None,         0),   # NULL статус - не включается
])
def test_filter_completed_various_statuses(spark, status, expected_count):
    """Проверяем, что фильтр обрабатывает различные статусы корректно."""
    df = spark.createDataFrame(
        [(1, status), (2, "completed")],  # строка 2 всегда completed
        ["id", "status"]
    )
    result = filter_completed_transactions(df)
    assert result.count() == expected_count + 1  # строка 2 всегда проходит

from pyspark.sql.types import IntegerType, LongType, DoubleType, FloatType

@pytest.mark.parametrize("amount_type", [
    IntegerType(),
    LongType(),
    DoubleType(),
    FloatType(),
])
def test_calculate_tax_works_with_numeric_types(spark, amount_type):
    """Функция должна работать с любыми числовыми типами для amount."""
    schema = StructType([
        StructField("id",     LongType(),   nullable=False),
        StructField("amount", amount_type,  nullable=True),
    ])
    df = spark.createDataFrame([(1, 100)], schema)
    # Тест не должен падать с TypeError или AnalysisException
    result = calculate_tax(df, tax_rate=0.2)
    assert "tax" in result.columns

@pytest.fixture(params=["completed", "COMPLETED", "Completed"])
def status_variant(request):
    """Различные варианты написания статуса 'completed'."""
    return request.param


def test_filter_is_case_sensitive(spark, status_variant):
    """Только точное совпадение 'completed' проходит фильтр."""
    df = spark.createDataFrame(
        [(1, status_variant)],
        ["id", "status"]
    )
    result = filter_completed_transactions(df)
    expected_count = 1 if status_variant == "completed" else 0
    assert result.count() == expected_count

# src/udfs.py
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType, DoubleType


def _classify_amount_impl(amount: float | None) -> str:
    """
    Чистая Python-функция - тестируется без Spark.
    Классифицирует транзакцию по сумме.
    """
    if amount is None:
        return "unknown"
    if amount < 0:
        return "refund"
    if amount == 0:
        return "zero"
    if amount < 100:
        return "small"
    if amount < 1000:
        return "medium"
    return "large"


# UDF - обёртка над чистой функцией
classify_amount_udf = udf(_classify_amount_impl, StringType())


def _normalize_phone_impl(phone: str | None) -> str | None:
    """Нормализует номер телефона к формату +7XXXXXXXXXX."""
    if phone is None:
        return None
    digits = "".join(c for c in phone if c.isdigit())
    if len(digits) == 11 and digits.startswith("8"):
        return "+7" + digits[1:]
    if len(digits) == 10:
        return "+7" + digits
    return None  # невалидный формат


normalize_phone_udf = udf(_normalize_phone_impl, StringType())

# tests/test_udfs.py
from src.udfs import _classify_amount_impl, _normalize_phone_impl


# Тест чистой функции - быстро, без JVM
@pytest.mark.parametrize("amount, expected", [
    (None,     "unknown"),
    (-50.0,    "refund"),
    (0.0,      "zero"),
    (50.0,     "small"),
    (500.0,    "medium"),
    (5000.0,   "large"),
    (99.99,    "small"),   # граничное значение
    (100.0,    "medium"),  # точно на границе
])
def test_classify_amount_logic(amount, expected):
    """Юнит-тест Python-логики без Spark - мгновенно выполняется."""
    assert _classify_amount_impl(amount) == expected


@pytest.mark.parametrize("phone, expected", [
    (None,            None),              # NULL входные данные
    ("+7 999 123-45-67",  "+79991234567"),  # уже в формате
    ("89991234567",   "+79991234567"),    # начинается с 8
    ("9991234567",    "+79991234567"),    # без кода страны
    ("123",           None),             # слишком короткий
    ("",              None),             # пустая строка
])
def test_normalize_phone_logic(phone, expected):
    assert _normalize_phone_impl(phone) == expected

from src.udfs import classify_amount_udf, normalize_phone_udf
from pyspark.sql import functions as F


def test_classify_amount_udf_in_dataframe(spark):
    """Проверяем, что UDF корректно интегрируется в DataFrame API."""
    df = spark.createDataFrame(
        [(1, -10.0), (2, 50.0), (3, 500.0), (4, None)],
        ["id", "amount"]
    )

    result = df.withColumn("category", classify_amount_udf(F.col("amount")))

    rows = {r["id"]: r["category"] for r in result.collect()}
    assert rows[1] == "refund"
    assert rows[2] == "small"
    assert rows[3] == "medium"
    assert rows[4] == "unknown"


def test_classify_udf_returns_string_type(spark):
    """Тип результата UDF должен быть StringType."""
    df = spark.createDataFrame([(1, 100.0)], ["id", "amount"])
    result = df.withColumn("category", classify_amount_udf(F.col("amount")))
    from pyspark.sql.types import StringType
    assert result.schema["category"].dataType == StringType()


def test_udf_handles_large_partition(spark):
    """UDF должен корректно обрабатывать большие партиции (не только 1-2 строки)."""
    import random
    data = [(i, random.uniform(-100, 10000)) for i in range(1000)]
    df = spark.createDataFrame(data, ["id", "amount"])

    result = df.withColumn("category", classify_amount_udf(F.col("amount")))
    # Не должно быть NULL в категории (если amount не NULL)
    null_count = result.filter(F.col("category").isNull()).count()
    assert null_count == 0

# src/udfs.py
from pyspark.sql.functions import pandas_udf
import pandas as pd


@pandas_udf("double")
def clean_amount_pandas_udf(amounts: pd.Series) -> pd.Series:
    """Pandas UDF: обнуляет отрицательные суммы и заполняет NULL нулями."""
    return amounts.fillna(0.0).clip(lower=0.0)

# tests/test_udfs.py
import pandas as pd

def test_clean_amount_logic():
    """Тест логики напрямую через pandas - без Spark."""
    series = pd.Series([None, -50.0, 0.0, 100.0, -0.01])
    result = series.fillna(0.0).clip(lower=0.0)

    assert result[0] == 0.0    # NULL → 0
    assert result[1] == 0.0    # отрицательное → 0
    assert result[4] == 0.0    # -0.01 → 0


def test_clean_amount_pandas_udf_in_spark(spark):
    """Интеграционный тест: Pandas UDF в DataFrame."""
    from src.udfs import clean_amount_pandas_udf
    from pyspark.sql import functions as F

    df = spark.createDataFrame(
        [(1, None), (2, -100.0), (3, 500.0)],
        ["id", "amount"]
    )
    result = df.withColumn("clean", clean_amount_pandas_udf(F.col("amount")))
    rows = {r["id"]: r["clean"] for r in result.collect()}

    assert rows[1] == 0.0
    assert rows[2] == 0.0
    assert rows[3] == 500.0

def test_calculate_tax_null_amount(spark):
    """NULL amount должен давать налог 0.0, а не NULL."""
    df = spark.createDataFrame(
        [(1, None), (2, 100.0)],
        ["id", "amount"]
    )

    result = calculate_tax(df)

    # Проверяем строку с NULL
    tax_for_null = (
        result.filter(F.col("id") == 1)
              .select("tax")
              .collect()[0]["tax"]
    )
    assert tax_for_null == 0.0, f"Ожидалось 0.0, получено {tax_for_null}"

    # NULL не должен «заразить» другие строки
    tax_for_normal = (
        result.filter(F.col("id") == 2)
              .select("tax")
              .collect()[0]["tax"]
    )
    assert abs(tax_for_normal - 20.0) < 0.001

from pyspark.sql.types import LongType


def test_filter_on_empty_dataframe(spark):
    """Пустой DataFrame на входе → пустой DataFrame на выходе, без ошибок."""
    schema = StructType([
        StructField("id",     LongType(),   nullable=True),
        StructField("status", StringType(), nullable=True),
        StructField("amount", DoubleType(), nullable=True),
    ])

    empty_df = spark.createDataFrame([], schema)
    result = filter_completed_transactions(empty_df)

    assert result.count() == 0
    assert result.schema == schema  # схема должна сохраниться

import pytest
from pyspark.sql.utils import AnalysisException


def test_calculate_tax_raises_on_missing_column(spark):
    """Функция должна выбросить исключение, если нет колонки amount."""
    df_without_amount = spark.createDataFrame(
        [(1, "completed")],
        ["id", "status"]
    )

    # Важно: вызов Action (count) запускает plan и раскрывает AnalysisException.
    # Сама трансформация - ленивая и не бросает ошибку при объявлении.
    with pytest.raises(Exception):
        calculate_tax(df_without_amount).count()

def test_calculate_tax_boundary_values(spark):
    """Граничные значения: ноль, отрицательное, большое число."""
    df = spark.createDataFrame(
        [
            (1,  0.0),           # нулевая сумма → налог 0
            (2, -100.0),         # возврат → налог отрицательный (корректно)
            (3,  1_000_000.0),   # большие суммы не должны переполняться
        ],
        ["id", "amount"]
    )

    result = calculate_tax(df, tax_rate=0.2)
    rows = {r["id"]: r["tax"] for r in result.collect()}

    assert rows[1] == 0.0
    assert rows[2] == -20.0           # возврат налога - ожидаемое поведение
    assert abs(rows[3] - 200_000.0) < 0.01

# tests/test_transformations.py
import pytest
from pyspark.sql import functions as F
from pyspark.sql.types import (
    StructType, StructField,
    LongType, StringType, DoubleType, BooleanType,
)
from chispa.dataframe_comparer import assert_df_equality

from src.transformations import (
    filter_completed_transactions,
    calculate_tax,
    add_discount_flag,
)


class TestFilterCompletedTransactions:

    def test_keeps_only_completed(self, spark):
        df = spark.createDataFrame(
            [(1, "completed"), (2, "pending"), (3, "completed"), (4, "cancelled")],
            ["id", "status"]
        )
        result = filter_completed_transactions(df)

        assert result.count() == 2
        statuses = {r["status"] for r in result.collect()}
        assert statuses == {"completed"}

    def test_empty_input_returns_empty(self, spark):
        schema = StructType([
            StructField("id",     LongType(),   nullable=True),
            StructField("status", StringType(), nullable=True),
        ])
        result = filter_completed_transactions(spark.createDataFrame([], schema))
        assert result.count() == 0

    def test_all_completed_keeps_all(self, spark):
        df = spark.createDataFrame(
            [(1, "completed"), (2, "completed")],
            ["id", "status"]
        )
        assert filter_completed_transactions(df).count() == 2


class TestCalculateTax:

    def test_standard_rate(self, spark):
        input_df = spark.createDataFrame(
            [(1, 100.0), (2, 250.0), (3, 50.0)],
            ["id", "amount"]
        )
        result = calculate_tax(input_df, tax_rate=0.2)
        expected = spark.createDataFrame(
            [(1, 100.0, 20.0), (2, 250.0, 50.0), (3, 50.0, 10.0)],
            ["id", "amount", "tax"]
        )
        assert_df_equality(result, expected, ignore_row_order=True)

    def test_null_amount_gives_zero_tax(self, spark):
        df = spark.createDataFrame([(1, None)], ["id", "amount"])
        tax = calculate_tax(df).collect()[0]["tax"]
        assert tax == 0.0, f"Ожидалось 0.0, получено {tax}"

    def test_zero_amount(self, spark):
        df = spark.createDataFrame([(1, 0.0)], ["id", "amount"])
        assert calculate_tax(df).collect()[0]["tax"] == 0.0

    def test_custom_tax_rate(self, spark):
        df = spark.createDataFrame([(1, 100.0)], ["id", "amount"])
        assert abs(calculate_tax(df, tax_rate=0.1).collect()[0]["tax"] - 10.0) < 0.001

    def test_result_schema_has_tax_column(self, spark):
        df = spark.createDataFrame([(1, 100.0)], ["id", "amount"])
        result = calculate_tax(df)
        assert "tax" in result.columns
        assert result.schema["tax"].dataType == DoubleType()


class TestAddDiscountFlag:

    def test_big_discount_flagged(self, spark):
        """Скидка > 50% суммы → big_discount = True."""
        df = spark.createDataFrame(
            [(1, 100.0, 60.0)],  # 60 > 50% от 100
            ["id", "amount", "discount"]
        )
        assert add_discount_flag(df).collect()[0]["big_discount"] is True

    def test_small_discount_not_flagged(self, spark):
        df = spark.createDataFrame(
            [(1, 100.0, 30.0)],  # 30 < 50% от 100
            ["id", "amount", "discount"]
        )
        assert add_discount_flag(df).collect()[0]["big_discount"] is False

    def test_null_discount_not_flagged(self, spark):
        """NULL скидка - не помечаем как большую."""
        df = spark.createDataFrame(
            [(1, 100.0, None)],
            ["id", "amount", "discount"]
        )
        assert add_discount_flag(df).collect()[0]["big_discount"] is False

    def test_zero_amount_not_flagged(self, spark):
        """Нулевая сумма - деление на ноль не происходит, флаг False."""
        df = spark.createDataFrame(
            [(1, 0.0, 10.0)],
            ["id", "amount", "discount"]
        )
        assert add_discount_flag(df).collect()[0]["big_discount"] is False

def test_inner_join_excludes_unmatched(spark):
    """Inner join не должен включать строки без совпадений."""
    transactions = spark.createDataFrame(
        [(1, "user_1", 100.0), (2, "user_999", 50.0)],  # user_999 не существует
        ["tx_id", "user_id", "amount"]
    )
    users = spark.createDataFrame(
        [("user_1", "Alice"), ("user_2", "Bob")],
        ["user_id", "name"]
    )

    result = transactions.join(users, "user_id", "inner")

    assert result.count() == 1
    assert result.collect()[0]["name"] == "Alice"


def test_left_join_keeps_all_transactions(spark):
    """Left join сохраняет все транзакции, даже без совпадения в users."""
    transactions = spark.createDataFrame(
        [(1, "user_1", 100.0), (2, "user_999", 50.0)],
        ["tx_id", "user_id", "amount"]
    )
    users = spark.createDataFrame(
        [("user_1", "Alice")],
        ["user_id", "name"]
    )

    result = transactions.join(users, "user_id", "left")

    assert result.count() == 2
    # Строка без совпадения должна иметь NULL в колонке name
    unknown = result.filter(F.col("user_id") == "user_999").collect()[0]
    assert unknown["name"] is None


def test_join_with_duplicate_keys_causes_cartesian(spark):
    """Если ключ дублируется в правой таблице - количество строк увеличивается."""
    transactions = spark.createDataFrame(
        [(1, "user_1", 100.0)],
        ["tx_id", "user_id", "amount"]
    )
    # Дубликат user_id в dimension-таблице - частая ошибка в реальных данных
    user_addresses = spark.createDataFrame(
        [("user_1", "Moscow"), ("user_1", "Saint Petersburg")],  # два адреса!
        ["user_id", "city"]
    )

    result = transactions.join(user_addresses, "user_id", "inner")

    # Ожидаем 2 строки - по одной на каждый адрес
    # Если тест упадёт - в dim-таблице нарушена уникальность ключей
    assert result.count() == 2

from pyspark.sql.window import Window


def rank_by_amount(df: DataFrame) -> DataFrame:
    window = Window.partitionBy("user_id").orderBy(F.col("amount").desc())
    return df.withColumn("rank", F.rank().over(window))


def test_rank_assigns_correct_positions(spark):
    df = spark.createDataFrame(
        [
            ("user_1", 300.0),
            ("user_1", 100.0),
            ("user_1", 200.0),
            ("user_2",  50.0),
        ],
        ["user_id", "amount"]
    )

    result = rank_by_amount(df)
    rows = {
        (r["user_id"], r["amount"]): r["rank"]
        for r in result.collect()
    }

    assert rows[("user_1", 300.0)] == 1
    assert rows[("user_1", 200.0)] == 2
    assert rows[("user_1", 100.0)] == 3
    assert rows[("user_2",  50.0)] == 1  # отдельный user - своя нумерация


def test_rank_handles_ties(spark):
    """При одинаковых суммах rank() даёт обеим строкам одинаковый ранг."""
    df = spark.createDataFrame(
        [("user_1", 100.0), ("user_1", 100.0)],
        ["user_id", "amount"]
    )

    result = rank_by_amount(df)
    ranks = sorted([r["rank"] for r in result.collect()])
    assert ranks == [1, 1]  # оба получают ранг 1 (не 1 и 2, как row_number())

def test_average_within_tolerance(spark):
    df = spark.createDataFrame(
        [(1, 10.0), (2, 20.0), (3, 30.0)],
        ["id", "value"]
    )
    result_avg = df.agg(F.avg("value").alias("avg_val")).collect()[0]["avg_val"]

    assert abs(result_avg - 20.0) < 0.001, \
        f"Ожидалось ~20.0, получено {result_avg}"

# Вместо прямого чтения из S3 внутри функции:
def load_and_enrich(spark, product_ids: list[int]) -> DataFrame:
    # Это нельзя легко мокировать
    dim_products = spark.read.parquet("s3a://bucket/dim/products/")
    transactions = spark.read.parquet("s3a://bucket/raw/transactions/")
    return transactions.join(dim_products, "product_id")

# Лучше: принимать DataFrame как аргументы
def enrich_transactions(
    transactions: DataFrame,
    dim_products: DataFrame,
) -> DataFrame:
    """Обогащает транзакции данными из справочника продуктов."""
    return transactions.join(dim_products, "product_id", "left")

# В тесте: создаём синтетические DataFrame напрямую
def test_enrich_transactions(spark):
    transactions = spark.createDataFrame(
        [(1, "P001", 100.0), (2, "P999", 50.0)],  # P999 не существует в справочнике
        ["tx_id", "product_id", "amount"]
    )
    dim_products = spark.createDataFrame(
        [("P001", "Widget A", "Electronics")],
        ["product_id", "name", "category"]
    )

    result = enrich_transactions(transactions, dim_products)

    # P999 должен получить NULL в колонках name и category
    p999 = result.filter(F.col("product_id") == "P999").collect()[0]
    assert p999["name"] is None
    assert p999["category"] is None

from unittest.mock import patch, MagicMock


def test_pipeline_reads_from_correct_path(spark):
    """Проверяем, что пайплайн читает данные из правильного пути."""
    mock_df = spark.createDataFrame(
        [(1, "completed", 100.0)],
        ["id", "status", "amount"]
    )

    # Мокируем spark.read.parquet чтобы вернуть синтетические данные
    with patch.object(spark.read, "parquet", return_value=mock_df) as mock_read:
        from src.pipeline import run
        run(spark, input_path="s3a://bucket/raw/", output_path="/tmp/output/")

        # Проверяем, что чтение было вызвано с правильным путём
        mock_read.assert_called_once_with("s3a://bucket/raw/")

def test_with_monkeypatch(spark, monkeypatch, tmp_path):
    """Тест, где чтение данных заменено на локальный файл."""
    # Создаём тестовые данные на локальном диске
    test_data = spark.createDataFrame(
        [(1, "completed", 100.0)],
        ["id", "status", "amount"]
    )
    test_data.write.parquet(str(tmp_path / "test_input"))

    # Заменяем путь к S3 на локальный путь
    monkeypatch.setenv("INPUT_PATH", str(tmp_path / "test_input"))
    monkeypatch.setenv("OUTPUT_PATH", str(tmp_path / "test_output"))

    from src import pipeline
    pipeline.run(spark)

    # Проверяем результат через чтение с диска
    result = spark.read.parquet(str(tmp_path / "test_output"))
    assert result.count() == 1

# tests/test_golden.py
import pytest
from pathlib import Path
from chispa.dataframe_comparer import assert_df_equality
from src.transformations import complex_loyalty_score


GOLDEN_PATH = Path(__file__).parent / "golden_data"


@pytest.fixture(scope="session")
def create_golden_data(spark):
    """
    Создаёт эталонные данные один раз.
    Запускать только при изменении эталона (флаг --update-golden).
    """
    return False  # False = не обновлять, True = обновить эталон


def test_loyalty_score_matches_golden(spark, create_golden_data):
    """Результат совпадает с эталонным датасетом."""
    input_df = spark.read.parquet(str(GOLDEN_PATH / "input"))

    result = complex_loyalty_score(input_df)

    if create_golden_data:
        # Режим обновления: записываем новый эталон
        result.write.mode("overwrite").parquet(str(GOLDEN_PATH / "expected"))
        pytest.skip("Golden data updated - re-run without --update-golden")

    expected = spark.read.parquet(str(GOLDEN_PATH / "expected"))
    assert_df_equality(result, expected, ignore_row_order=True)

# Пример концепции (псевдокод - Deequ работает через Scala API)
# В Python доступен PyDeequ

from pydeequ.checks import Check, CheckLevel
from pydeequ.verification import VerificationSuite

check = (
    Check(spark, CheckLevel.Error, "Transaction Data Quality")
    .hasSize(lambda size: size > 0)                           # нет пустых датасетов
    .isComplete("order_id")                                   # order_id никогда не NULL
    .isUnique("order_id")                                     # order_id уникален
    .isNonNegative("amount")                                  # суммы неотрицательны
    .isContainedIn("status", ["completed", "pending", "cancelled"])  # только допустимые статусы
    .hasCompleteness("customer_id", lambda r: r >= 0.99)     # customer_id заполнен в 99%+ строк
)

result = VerificationSuite(spark).onData(df).addCheck(check).run()

# Great Expectations в production pipeline
import great_expectations as ge

df_ge = ge.dataset.SparkDFDataset(spark_df)

result = df_ge.expect_column_values_to_not_be_null("order_id")
assert result["success"], f"order_id содержит NULL: {result}"

result = df_ge.expect_column_values_to_be_between("amount", min_value=0)
assert result["success"], f"Найдены отрицательные суммы: {result}"

# 1. Чтение реальных файлов в unit-тестах
def test_reads_from_s3(spark):
    df = spark.read.parquet("s3a://prod-bucket/data/")
    # Тест зависит от состояния S3, сети и прав доступа - это не unit-тест

# 2. collect() на потенциально больших данных
def test_collects_everything(spark):
    result = run_heavy_pipeline(spark)
    all_data = result.collect()  # если 100M строк - Driver упадёт

# 3. Чувствительность к порядку строк
def test_order_sensitive(spark):
    result = transform(df)
    expected = [Row(id=1, val="a"), Row(id=2, val="b")]
    assert result.collect() == expected  # порядок не гарантирован!

# 4. Слишком много логики в одном тесте
def test_does_everything(spark):
    # Если тест упадёт - непонятно, что именно сломалось
    result = big_pipeline(df)
    assert result.count() == 5
    assert result.filter(F.col("tax").isNull()).count() == 0
    assert "big_discount" in result.columns
    assert result.schema["amount"].dataType == DoubleType()

# 5. Нет явной схемы в тестовых данных
def test_no_schema(spark):
    df = spark.createDataFrame([(1, 100)])
    # id: LongType, amount: LongType - но функция ожидает DoubleType!

# 6. Тестирование Spark internals вместо бизнес-логики
def test_spark_internals(spark):
    df = spark.createDataFrame([(1, 100.0)], ["id", "amount"])
    assert df.rdd.getNumPartitions() == 1  # это деталь реализации Spark, не бизнес-логика

# .github/workflows/test.yml
name: PySpark Unit Tests

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

jobs:
  test:
    runs-on: ubuntu-latest

    steps:
      - uses: actions/checkout@v4

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

      - name: Set up Java
        uses: actions/setup-java@v4
        with:
          distribution: temurin
          java-version: "17"

      - name: Cache pip
        uses: actions/cache@v4
        with:
          path: ~/.cache/pip
          key: ${{ runner.os }}-pip-${{ hashFiles('requirements-dev.txt') }}

      - name: Install dependencies
        run: pip install -r requirements-dev.txt

      - name: Run unit tests with coverage
        run: |
          pytest tests/ -v --tb=short --no-header \
            --cov=src \
            --cov-report=term-missing \
            --cov-fail-under=80
        env:
          PYSPARK_PYTHON: python3

      - name: Upload coverage report
        uses: codecov/codecov-action@v4
        if: always()

# Dockerfile.test
FROM python:3.11-slim

# Установка Java
RUN apt-get update && apt-get install -y --no-install-recommends \
    openjdk-17-jre-headless \
    && rm -rf /var/lib/apt/lists/*

ENV JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64

WORKDIR /app

COPY requirements-dev.txt .
RUN pip install --no-cache-dir -r requirements-dev.txt

COPY src/ src/
COPY tests/ tests/
COPY conftest.py .

CMD ["pytest", "tests/", "-v", "--tb=short", "--cov=src", "--cov-report=term-missing"]

# Сборка и запуск тестов в Docker
docker build -f Dockerfile.test -t pyspark-tests .
docker run --rm pyspark-tests

# Высокий coverage, но плохие тесты:
def test_calculate_tax_coverage_hack(spark):
    df = spark.createDataFrame([(1, 100.0)], ["id", "amount"])
    result = calculate_tax(df)
    assert result is not None  # выполнили всю функцию, но ничего не проверили!

pip install pytest chispa pyspark

# transformations.py
from pyspark.sql import DataFrame
from pyspark.sql import functions as F


def calculate_bonus(df: DataFrame, threshold: float = 1000.0) -> DataFrame:
    return df.withColumn(
        "bonus",
        F.when(F.col("amount") > threshold, F.col("amount") * 0.1)
         .otherwise(F.lit(0.0))
    )

# conftest.py
import pytest
from pyspark.sql import SparkSession


@pytest.fixture(scope="session")
def spark():
    s = (
        SparkSession.builder
        .master("local[*]")
        .appName("tests")
        .config("spark.sql.shuffle.partitions", "4")
        .config("spark.ui.enabled", "false")
        .getOrCreate()
    )
    yield s
    s.stop()

# test_transformations.py
import pytest
from chispa.dataframe_comparer import assert_df_equality
from transformations import calculate_bonus


def test_bonus_above_threshold(spark):
    df = spark.createDataFrame([(1, 1500.0), (2, 500.0)], ["id", "amount"])
    result = calculate_bonus(df, threshold=1000.0)
    expected = spark.createDataFrame(
        [(1, 1500.0, 150.0), (2, 500.0, 0.0)],
        ["id", "amount", "bonus"]
    )
    assert_df_equality(result, expected, ignore_row_order=True)


def test_null_amount_gives_zero_bonus(spark):
    df = spark.createDataFrame([(1, None)], ["id", "amount"])
    result = calculate_bonus(df)
    assert result.collect()[0]["bonus"] == 0.0


@pytest.mark.parametrize("amount, expected_bonus", [
    (0.0,    0.0),    # ноль
    (999.9,  0.0),    # ниже порога
    (1000.0, 0.0),    # ровно на пороге (не включительно)
    (1000.1, 100.01), # выше порога
    (5000.0, 500.0),  # большая сумма
])
def test_bonus_threshold_boundary(spark, amount, expected_bonus):
    df = spark.createDataFrame([(1, amount)], ["id", "amount"])
    result = calculate_bonus(df, threshold=1000.0)
    actual = result.collect()[0]["bonus"]
    assert abs(actual - expected_bonus) < 0.001

pytest -v
# test_bonus_above_threshold         PASSED
# test_null_amount_gives_zero_bonus  PASSED
# test_bonus_threshold_boundary[0.0-0.0]    PASSED
# test_bonus_threshold_boundary[999.9-0.0]  PASSED
# test_bonus_threshold_boundary[1000.0-0.0] PASSED
# test_bonus_threshold_boundary[1000.1-100.01] PASSED
# test_bonus_threshold_boundary[5000.0-500.0]  PASSED