Testing PySpark: pyspark.testing.assertDataFrameEqual и Unit-тесты

# АНТИПАТТЕРН: нетестируемый монолитный ETL
# Весь код - процедурная последовательность, IO и трансформации перемешаны

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

# SparkSession создаётся прямо в скрипте - невозможно подменить в тесте
spark = SparkSession.builder.appName("etl").getOrCreate()

# IO и бизнес-логика перемешаны - нельзя тестировать логику без чтения данных
df = spark.read.parquet("s3://bronze/orders/")

# Сложная трансформация прямо в main-скрипте - никакого переиспользования
df_clean = df \
    .filter(F.col("order_id").isNotNull()) \
    .withColumn("total", F.col("quantity") * F.col("unit_price")) \
    .withColumn("total", F.when(F.col("discount") > 0,
                                F.col("total") * (1 - F.col("discount")))
                          .otherwise(F.col("total")))

df_enriched = df_clean \
    .join(spark.read.parquet("s3://silver/customers/"),  # ещё один IO внутри логики!
          on="customer_id", how="left") \
    .dropDuplicates(["order_id"])

df_enriched.write.mode("overwrite").parquet("s3://silver/orders/")
spark.stop()
# КАК ТЕСТИРОВАТЬ ЭТО? Нужны s3://bronze/, s3://silver/, реальный Spark...

"""
orders_pipeline.py
Правильная архитектура: IO отделён от логики.
Трансформационные функции принимают DataFrame и возвращают DataFrame.
"""

from pyspark.sql import DataFrame, functions as F
from pyspark.sql.types import StructType, StructField, StringType, DoubleType, IntegerType


# ── Transformation Layer (чистые функции - легко тестировать) ─────────────────

def clean_orders(df: DataFrame) -> DataFrame:
    """
    Очищает сырые заказы:
    - Убирает строки с NULL order_id (невалидные записи)
    - Приводит типы к стандарту
    - Убирает leading/trailing spaces в строковых полях

    Чистая функция: только DataFrame → DataFrame, нет IO, нет side effects.
    Именно такие функции легко unit-тестировать.
    """
    return (
        df
        .filter(F.col("order_id").isNotNull())
        .filter(F.col("quantity") > 0)
        .withColumn("order_id", F.trim(F.col("order_id")))
        .withColumn("quantity",   F.col("quantity").cast(IntegerType()))
        .withColumn("unit_price", F.col("unit_price").cast(DoubleType()))
        .withColumn("discount",   F.coalesce(F.col("discount"), F.lit(0.0)))
    )


def calculate_order_total(df: DataFrame) -> DataFrame:
    """
    Вычисляет итоговую сумму заказа с учётом скидки.

    Формула: total = quantity * unit_price * (1 - discount)
    Скидка: от 0.0 (0%) до 1.0 (100%).

    Граничный случай: discount > 1.0 → total становится отрицательным.
    Это баг, который unit-тест должен поймать.
    """
    return df.withColumn(
        "order_total",
        F.round(
            F.col("quantity") * F.col("unit_price") * (1 - F.col("discount")),
            2
        )
    )


def enrich_with_customer(orders: DataFrame, customers: DataFrame) -> DataFrame:
    """
    Обогащает заказы данными клиента (имя, сегмент).

    LEFT JOIN: заказы без клиента остаются в результате с NULL в customer_name.
    ВАЖНО: деревянная ловушка - если customers содержит дубли customer_id,
    JOIN умножит строки заказов. Это должен поймать unit-тест.
    """
    customers_deduped = customers.dropDuplicates(["customer_id"])
    return orders.join(
        customers_deduped.select("customer_id", "customer_name", "segment"),
        on="customer_id",
        how="left",
    )


def deduplicate_orders(df: DataFrame) -> DataFrame:
    """
    Дедупликация по order_id: при дублях оставляем самую свежую запись.

    Используем Window + row_number вместо dropDuplicates,
    чтобы контролировать, какую из дублирующихся строк оставить.
    """
    from pyspark.sql.window import Window

    window_spec = Window.partitionBy("order_id").orderBy(
        F.col("updated_at").desc_nulls_last()
    )
    return (
        df
        .withColumn("_rn", F.row_number().over(window_spec))
        .filter(F.col("_rn") == 1)
        .drop("_rn")
    )


# ── IO Layer (тестируем только integration-тестами) ───────────────────────────

def read_bronze_orders(spark, date: str) -> DataFrame:
    """Читает сырые заказы из Bronze-слоя."""
    return spark.read.parquet(f"s3://bronze/orders/date={date}/")


def write_silver_orders(df: DataFrame, target: str, date: str) -> None:
    """Записывает обработанные заказы в Silver-слой."""
    (
        df.write
        .format("iceberg")
        .mode("overwrite")
        .option("replaceWhere", f"order_date = '{date}'")
        .saveAsTable(target)
    )


# ── Main (оркестрирует IO + transforms) ──────────────────────────────────────

def run_pipeline(spark, date: str, target: str = "catalog.silver.orders") -> None:
    """Основной пайплайн. IO + оркестрация трансформаций."""
    customers = spark.read.parquet("s3://silver/customers/")

    orders_raw = read_bronze_orders(spark, date)

    orders_clean = clean_orders(orders_raw)
    orders_total = calculate_order_total(orders_clean)
    orders_enriched = enrich_with_customer(orders_total, customers)
    orders_final = deduplicate_orders(orders_enriched)

    write_silver_orders(orders_final, target, date)

# tests/conftest.py
# Этот файл автоматически подгружается pytest для всех тестов в директории.
# Содержит разделяемые fixtures.

import pytest
from pyspark.sql import SparkSession


@pytest.fixture(scope="session")
def spark():
    """
    Разделяемая SparkSession для всего тест-сьюта.

    scope="session": создаётся ОДИН РАЗ для всего тест-запуска.
    Это критично для скорости: создание SparkSession занимает ~5-10 секунд.

    Конфигурация минимальная: local[2] для использования 2 потоков,
    малый объём памяти (тесты работают с маленькими датасетами).
    """
    spark_session = (
        SparkSession.builder
        .master("local[2]")                          # 2 потока - достаточно для тестов
        .appName("pyspark-unit-tests")
        .config("spark.driver.memory", "2g")         # Минимум памяти (тестовые данные маленькие)
        .config("spark.sql.shuffle.partitions", "4") # Меньше партиций → быстрее для малых данных
        .config("spark.sql.adaptive.enabled", "false") # Отключаем AQE: тесты должны быть детерминированными
        .config("spark.ui.showConsoleProgress", "false") # Меньше шума в логах тестов
        .config("spark.ui.enabled", "false")         # Не нужен Spark UI в тестах
        # Отключаем логи Spark: в тестах важен вывод pytest, не Spark INFO
        .config("spark.sql.execution.arrow.pyspark.enabled", "false")
        .getOrCreate()
    )

    # Снижаем уровень логов до WARN: меньше шума
    spark_session.sparkContext.setLogLevel("WARN")

    yield spark_session  # Передаём сессию тестам

    # Teardown: вызывается автоматически после последнего теста
    spark_session.stop()


@pytest.fixture(scope="function")
def sample_orders(spark):
    """
    Минимальный тестовый датасет заказов.

    scope="function": создаётся заново для каждого теста.
    Это гарантирует изоляцию: один тест не может испортить данные для другого.

    Правило хорошего теста: 3-10 строк с явными, читаемыми значениями.
    Каждая строка должна тестировать что-то конкретное.
    """
    from pyspark.sql.types import StructType, StructField, StringType, IntegerType, DoubleType

    schema = StructType([
        StructField("order_id",    StringType(),  nullable=True),
        StructField("customer_id", StringType(),  nullable=False),
        StructField("quantity",    IntegerType(), nullable=True),
        StructField("unit_price",  DoubleType(),  nullable=True),
        StructField("discount",    DoubleType(),  nullable=True),
        StructField("updated_at",  StringType(),  nullable=True),
    ])

    data = [
        ("ORD-001", "CUST-01", 2, 10.0, 0.1, "2024-01-15 10:00:00"),  # Нормальный заказ
        ("ORD-002", "CUST-02", 1, 50.0, 0.0, "2024-01-15 11:00:00"),  # Без скидки
        (None,      "CUST-03", 3, 20.0, 0.0, "2024-01-15 12:00:00"),  # NULL order_id → должен быть отфильтрован
        ("ORD-003", "CUST-01", 0, 15.0, 0.0, "2024-01-15 13:00:00"),  # qty=0 → должен быть отфильтрован
        ("ORD-001", "CUST-01", 2, 10.0, 0.1, "2024-01-15 14:00:00"),  # Дубль ORD-001 (более свежий)
    ]

    return spark.createDataFrame(data, schema=schema)


@pytest.fixture(scope="function")
def sample_customers(spark):
    """Тестовые данные клиентов для join-тестов."""
    from pyspark.sql.types import StructType, StructField, StringType

    schema = StructType([
        StructField("customer_id",   StringType(), nullable=False),
        StructField("customer_name", StringType(), nullable=True),
        StructField("segment",       StringType(), nullable=True),
    ])

    data = [
        ("CUST-01", "Alice Smith", "Premium"),
        ("CUST-02", "Bob Johnson", "Standard"),
        # CUST-03 отсутствует - тест на LEFT JOIN behaviour
    ]

    return spark.createDataFrame(data, schema=schema)

project/
├── src/
│   └── orders_pipeline.py        ← Код с чистыми transform-функциями
├── tests/
│   ├── conftest.py               ← SparkSession fixture и shared данные
│   ├── unit/
│   │   ├── test_clean_orders.py  ← Unit-тесты для clean_orders()
│   │   ├── test_calculate_total.py
│   │   ├── test_enrich.py
│   │   └── test_dedup.py
│   ├── integration/
│   │   └── test_pipeline_e2e.py  ← Integration-тесты (с реальным хранилищем)
│   └── schema/
│       └── test_schema_contracts.py  ← Тесты схемы данных
├── pyproject.toml
└── pytest.ini

# pytest.ini
[pytest]
testpaths = tests
python_files = test_*.py
python_classes = Test*
python_functions = test_*
addopts =
    -v                    # Verbose: показывать имена тестов
    --tb=short            # Короткий traceback при падении
    --no-header           # Без лишнего заголовка
filterwarnings =
    ignore::DeprecationWarning
    ignore::UserWarning:pyspark

# АНТИПАТТЕРН: наивное сравнение через collect()
def test_clean_bad():
    result = clean_orders(sample_df)
    expected = [("ORD-001", "CUST-01", ...), ...]

    # Проблема 1: порядок строк в Spark НЕ ГАРАНТИРОВАН
    # Distributed execution → разные партиции → разный порядок при collect()
    # Тест будет мигать (flaky): иногда проходит, иногда нет
    assert result.collect() == expected  # Nondeterministic!

    # Проблема 2: схема не проверяется
    # IntegerType vs LongType → collect() вернёт одинаковые Python int
    # Но в Spark это разные типы, и это важно для downstream

    # Проблема 3: float precision
    # 18.0 == 18.000000001 → False из-за floating point
    # Тест упадёт на корректных данных

from pyspark.testing import assertDataFrameEqual
from pyspark.testing.utils import assertSchemaEqual
from pyspark.sql.types import StructType, StructField, StringType, IntegerType, DoubleType

def test_assert_dataframe_equal_anatomy(spark):
    """
    Демонстрация ключевых параметров assertDataFrameEqual.
    """

    # Сравниваемые DataFrames
    result = spark.createDataFrame(
        [("ORD-002", "CUST-02", 1, 50.0, 0.0, 50.0),
         ("ORD-001", "CUST-01", 2, 10.0, 0.1, 18.0)],
        schema="order_id STRING, customer_id STRING, quantity INT, "
               "unit_price DOUBLE, discount DOUBLE, order_total DOUBLE"
    )

    expected = spark.createDataFrame(
        [("ORD-001", "CUST-01", 2, 10.0, 0.1, 18.0),
         ("ORD-002", "CUST-02", 1, 50.0, 0.0, 50.0)],
        schema="order_id STRING, customer_id STRING, quantity INT, "
               "unit_price DOUBLE, discount DOUBLE, order_total DOUBLE"
    )

    assertDataFrameEqual(
        actual=result,
        expected=expected,

        # checkRowOrder=False - КРИТИЧНО для Spark:
        # Distributed execution не гарантирует порядок строк.
        # С False: строки сортируются перед сравнением → детерминированный результат.
        # С True (default): тест будет мигать (flaky) на разных кластерах.
        checkRowOrder=False,

        # checkSchema=True (default) - сравнивать схему:
        # Типы данных (INT vs LONG), nullable флаги, имена колонок.
        # Важно! IntegerType vs LongType - это разные типы в Spark,
        # хотя Python-int одинаковый.
        checkSchema=True,

        # rtol=0.0 (default) - относительная погрешность для float сравнения.
        # rtol=1e-5 означает: |actual - expected| / |expected| <= 1e-5
        # Используйте для финансовых расчётов: rtol=1e-10
        rtol=0.0,

        # atol=0.0 (default) - абсолютная погрешность.
        # atol=0.01 означает: |actual - expected| <= 0.01
        # Используйте когда expected ≈ 0 (rtol не работает при expected=0)
        atol=0.0,
    )
    # Если DataFrames совпадают → тест проходит без исключений
    # Если нет → AssertionError с понятным diff-сообщением

# Демонстрация проблемы без checkRowOrder=False
def test_row_order_problem(spark):
    """
    Показывает, почему checkRowOrder=True (default) опасен.
    """
    # Spark может вернуть строки в ЛЮБОМ порядке после shuffle
    df = spark.range(5).orderBy(F.rand())  # Случайный порядок

    # Этот тест будет мигать:
    # Иногда порядок совпадёт, иногда нет
    try:
        assertDataFrameEqual(df, spark.range(5), checkRowOrder=True)
    except AssertionError:
        print("Порядок строк не совпал - это ОЖИДАЕМО для Spark!")

    # Этот тест стабилен:
    # Строки сортируются перед сравнением
    assertDataFrameEqual(df, spark.range(5), checkRowOrder=False)  # Всегда проходит

def test_float_precision(spark):
    """
    Демонстрация управления точностью float сравнений.

    Проблема: 2 * 10.0 * (1 - 0.1) = 18.0 в теории,
    но 18.000000000000004 на практике из-за IEEE 754.
    """
    result = spark.createDataFrame(
        [("ORD-001", 18.000000000000004)],
        schema="order_id STRING, order_total DOUBLE"
    )
    expected = spark.createDataFrame(
        [("ORD-001", 18.0)],
        schema="order_id STRING, order_total DOUBLE"
    )

    # БЕЗ atol: тест упадёт (18.0 != 18.000000000000004)
    # assertDataFrameEqual(result, expected)  # AssertionError!

    # С atol=1e-6: допускаем абсолютное отклонение до 0.000001
    assertDataFrameEqual(result, expected, checkRowOrder=False, atol=1e-6)  # Проходит

    # С rtol=1e-6: допускаем относительное отклонение до 0.0001%
    assertDataFrameEqual(result, expected, checkRowOrder=False, rtol=1e-6)  # Проходит

def test_schema_relaxed(spark):
    """
    checkSchema=False: игнорируем типы данных, проверяем только значения.

    Полезно когда источник присылает INT, а мы ожидаем LONG,
    но значения одинаковые. Или когда nullable-флаги разные.
    """
    result = spark.createDataFrame(
        [("ORD-001", 2)],   # INT в результате
        schema="order_id STRING, quantity INT"
    )
    expected = spark.createDataFrame(
        [("ORD-001", 2)],   # LONG в expected (spark.createDataFrame default для int)
        schema="order_id STRING, quantity LONG"
    )

    # С checkSchema=True: AssertionError (INT != LONG)
    # С checkSchema=False: проходит (значения одинаковые)
    assertDataFrameEqual(result, expected, checkRowOrder=False, checkSchema=False)

from pyspark.testing.utils import assertSchemaEqual
from pyspark.sql.types import StructType, StructField, StringType, IntegerType, DoubleType

def test_output_schema(spark):
    """
    Проверяем, что функция возвращает DataFrame с правильной схемой,
    независимо от значений данных.

    Это особенно важно после изменений: добавили колонку? Изменили тип?
    Тест немедленно это поймает.
    """
    sample_df = spark.createDataFrame(
        [("ORD-001", "CUST-01", 2, 10.0, 0.1, None)],
        schema="order_id STRING, customer_id STRING, quantity INT, "
               "unit_price DOUBLE, discount DOUBLE, updated_at STRING"
    )

    result = calculate_order_total(sample_df)

    expected_schema = StructType([
        StructField("order_id",    StringType(),  True),
        StructField("customer_id", StringType(),  False),
        StructField("quantity",    IntegerType(), True),
        StructField("unit_price",  DoubleType(),  True),
        StructField("discount",    DoubleType(),  True),
        StructField("updated_at",  StringType(),  True),
        StructField("order_total", DoubleType(),  True),  # Новая колонка
    ])

    assertSchemaEqual(
        actual=result.schema,
        expected=expected_schema,
        # checkNullable=True (default): проверяет nullable флаги
        # checkNullable=False: игнорирует nullable, проверяет только имена и типы
        checkNullable=False,  # Часто нужно False: nullable зависит от операции
    )

# tests/unit/test_clean_orders.py

import pytest
from pyspark.testing import assertDataFrameEqual
from src.orders_pipeline import clean_orders


class TestCleanOrders:
    """
    Тесты для функции clean_orders().

    Каждый тест проверяет ОДИН конкретный аспект поведения.
    Имена тестов описывают сценарий: test_{condition}_{expected_behaviour}.
    """

    def test_removes_null_order_id(self, spark, sample_orders):
        """
        Строки с NULL order_id должны быть удалены.
        Это нарушение Data Contract: order_id - обязательное поле.
        """
        result = clean_orders(sample_orders)
        null_count = result.filter("order_id IS NULL").count()
        assert null_count == 0, "NULL order_id rows should be removed"

    def test_removes_zero_quantity(self, spark, sample_orders):
        """
        Строки с quantity=0 должны быть удалены.
        Заказ с нулевым количеством - это ошибка источника.
        """
        result = clean_orders(sample_orders)
        zero_qty_count = result.filter("quantity = 0").count()
        assert zero_qty_count == 0, "Zero quantity orders should be removed"

    def test_fills_null_discount_with_zero(self, spark, spark_session=None):
        """
        NULL в поле discount должен заменяться на 0.0 (нет скидки).
        Это бизнес-правило: отсутствие скидки != NULL.
        """
        spark_session = spark
        df = spark_session.createDataFrame(
            [("ORD-001", "CUST-01", 2, 10.0, None, "2024-01-15")],
            schema="order_id STRING, customer_id STRING, quantity INT, "
                   "unit_price DOUBLE, discount DOUBLE, updated_at STRING"
        )

        result = clean_orders(df)
        discount_val = result.select("discount").collect()[0]["discount"]
        assert discount_val == 0.0, "NULL discount should be filled with 0.0"

    def test_trims_whitespace_in_order_id(self, spark):
        """
        Пробелы в order_id должны быть удалены.
        Источник иногда присылает ' ORD-001 ' вместо 'ORD-001'.
        """
        df = spark.createDataFrame(
            [("  ORD-001  ", "CUST-01", 1, 10.0, 0.0, "2024-01-15")],
            schema="order_id STRING, customer_id STRING, quantity INT, "
                   "unit_price DOUBLE, discount DOUBLE, updated_at STRING"
        )

        result = clean_orders(df)
        order_id = result.select("order_id").collect()[0]["order_id"]
        assert order_id == "ORD-001", f"Expected 'ORD-001', got '{order_id}'"

    def test_returns_correct_rows_from_mixed_input(self, spark, sample_orders):
        """
        Из смешанного ввода (валидные + невалидные строки) должны
        вернуться только валидные.

        sample_orders содержит:
        - 2 валидных строки (ORD-001 × 2 дубля, ORD-002)
        - 1 строку с NULL order_id → отфильтровать
        - 1 строку с quantity=0 → отфильтровать
        Ожидаем 3 строки: ORD-001 (два дубля остаются!), ORD-002
        """
        result = clean_orders(sample_orders)
        assert result.count() == 3, f"Expected 3 rows, got {result.count()}"

# tests/unit/test_calculate_total.py

from pyspark.testing import assertDataFrameEqual
from src.orders_pipeline import calculate_order_total


def test_calculate_total_with_discount(spark):
    """
    Проверяем формулу: total = qty * price * (1 - discount)
    Тест с конкретными числами, которые легко проверить вручную.
    """
    from pyspark.sql import Row

    input_df = spark.createDataFrame([
        Row(order_id="ORD-001", quantity=2, unit_price=10.0, discount=0.1),
        Row(order_id="ORD-002", quantity=1, unit_price=50.0, discount=0.0),
        Row(order_id="ORD-003", quantity=3, unit_price=20.0, discount=0.5),
    ])

    result = calculate_order_total(input_df)

    expected = spark.createDataFrame([
        Row(order_id="ORD-001", quantity=2, unit_price=10.0, discount=0.1, order_total=18.0),
        Row(order_id="ORD-002", quantity=1, unit_price=50.0, discount=0.0, order_total=50.0),
        Row(order_id="ORD-003", quantity=3, unit_price=20.0, discount=0.5, order_total=30.0),
    ])

    assertDataFrameEqual(
        result.select("order_id", "order_total"),
        expected.select("order_id", "order_total"),
        checkRowOrder=False,
        atol=1e-6,  # Допускаем floating point погрешность
    )


def test_calculate_total_full_discount(spark):
    """
    Граничный случай: скидка 100% (discount=1.0) → total = 0.
    """
    input_df = spark.createDataFrame(
        [("ORD-001", 5, 100.0, 1.0)],
        schema="order_id STRING, quantity INT, unit_price DOUBLE, discount DOUBLE"
    )
    result = calculate_order_total(input_df)
    total = result.select("order_total").collect()[0]["order_total"]
    assert total == 0.0, f"100% discount should give 0 total, got {total}"


def test_calculate_total_preserves_other_columns(spark):
    """
    calculate_order_total добавляет колонку, но не должна менять другие.
    Проверяем через assertDataFrameEqual с checkSchema=True.
    """
    input_df = spark.createDataFrame(
        [("ORD-001", "CUST-01", 2, 10.0, 0.0)],
        schema="order_id STRING, customer_id STRING, quantity INT, unit_price DOUBLE, discount DOUBLE"
    )
    result = calculate_order_total(input_df)

    # Все исходные колонки должны присутствовать без изменений
    assert "customer_id" in result.columns
    assert "quantity" in result.columns
    # Новая колонка добавлена
    assert "order_total" in result.columns

# tests/unit/test_enrich.py

from pyspark.testing import assertDataFrameEqual
from src.orders_pipeline import enrich_with_customer


def test_enrich_matches_known_customer(spark):
    """
    Основной happy path: клиент есть в customers → поля заполнены.
    """
    orders = spark.createDataFrame(
        [("ORD-001", "CUST-01", 18.0)],
        schema="order_id STRING, customer_id STRING, order_total DOUBLE"
    )
    customers = spark.createDataFrame(
        [("CUST-01", "Alice Smith", "Premium")],
        schema="customer_id STRING, customer_name STRING, segment STRING"
    )

    result = enrich_with_customer(orders, customers)
    row = result.collect()[0]

    assert row["customer_name"] == "Alice Smith"
    assert row["segment"] == "Premium"


def test_enrich_preserves_order_without_customer(spark):
    """
    LEFT JOIN: заказ без совпадающего клиента остаётся в результате.
    customer_name и segment → NULL (это ожидаемое поведение).

    ВАЖНО: это критичный тест! Если бы использовали INNER JOIN —
    заказы без клиентов потерялись бы незаметно.
    """
    orders = spark.createDataFrame(
        [("ORD-999", "UNKNOWN-CUST", 50.0)],
        schema="order_id STRING, customer_id STRING, order_total DOUBLE"
    )
    customers = spark.createDataFrame(
        [("CUST-01", "Alice Smith", "Premium")],
        schema="customer_id STRING, customer_name STRING, segment STRING"
    )

    result = enrich_with_customer(orders, customers)
    assert result.count() == 1, "Order without customer should remain"

    row = result.collect()[0]
    assert row["customer_name"] is None, "customer_name should be NULL for unknown customer"


def test_enrich_deduplicates_customers_before_join(spark):
    """
    ЛОВУШКА: если customers содержит дубли customer_id,
    LEFT JOIN размножит строки заказов.

    Функция enrich_with_customer ДОЛЖНА дедуплицировать customers
    перед JOIN. Этот тест проверяет, что дедупликация работает.

    Это реальный production-баг, который встречается часто:
    'customers' пришёл без dedup → в Gold появились дубли заказов.
    """
    orders = spark.createDataFrame(
        [("ORD-001", "CUST-01", 18.0)],
        schema="order_id STRING, customer_id STRING, order_total DOUBLE"
    )
    customers_with_dupes = spark.createDataFrame(
        [("CUST-01", "Alice Smith", "Premium"),
         ("CUST-01", "Alice Smith", "Premium")],  # Дубль!
        schema="customer_id STRING, customer_name STRING, segment STRING"
    )

    result = enrich_with_customer(orders, customers_with_dupes)

    # БЕЗ dedup в функции: count() = 2 (JOIN умножил строку заказа на 2)
    # С dedup: count() = 1
    assert result.count() == 1, \
        "Duplicate customers should not multiply orders. " \
        "Make sure enrich_with_customer deduplicates customers before JOIN."

# tests/unit/test_dedup.py

from pyspark.testing import assertDataFrameEqual
from src.orders_pipeline import deduplicate_orders


def test_dedup_removes_duplicates(spark):
    """
    Из двух строк с одинаковым order_id должна остаться одна.
    """
    df = spark.createDataFrame(
        [("ORD-001", "2024-01-15 10:00:00", 18.0),
         ("ORD-001", "2024-01-15 14:00:00", 18.0)],  # Дубль
        schema="order_id STRING, updated_at STRING, order_total DOUBLE"
    )
    result = deduplicate_orders(df)
    assert result.count() == 1, "Duplicate order_id should result in 1 row"


def test_dedup_keeps_latest_by_updated_at(spark):
    """
    При дублях должна остаться строка с максимальным updated_at.

    Это бизнес-правило: более свежая версия заказа перекрывает старую.
    """
    df = spark.createDataFrame(
        [("ORD-001", "2024-01-15 10:00:00", 15.0),   # Старая версия
         ("ORD-001", "2024-01-15 14:00:00", 18.0)],  # Новая версия (скидка изменилась)
        schema="order_id STRING, updated_at STRING, order_total DOUBLE"
    )
    result = deduplicate_orders(df)

    row = result.collect()[0]
    assert row["updated_at"] == "2024-01-15 14:00:00", \
        "Should keep the most recent version (max updated_at)"
    assert row["order_total"] == 18.0, \
        "order_total should match the most recent version"


def test_dedup_handles_null_updated_at(spark):
    """
    Граничный случай: updated_at = NULL.
    Строки с NULL updated_at должны идти ПОСЛЕ строк с реальным timestamp.

    В deduplicate_orders используем orderBy(...desc_nulls_last()),
    что означает: NULL → последние в Window → удаляются при row_number > 1.
    """
    df = spark.createDataFrame(
        [("ORD-001", None, 15.0),                    # NULL updated_at
         ("ORD-001", "2024-01-15 10:00:00", 18.0)],  # Реальный timestamp
        schema="order_id STRING, updated_at STRING, order_total DOUBLE"
    )
    result = deduplicate_orders(df)
    assert result.count() == 1
    row = result.collect()[0]
    # Должна остаться строка с реальным timestamp, не NULL
    assert row["updated_at"] == "2024-01-15 10:00:00"


def test_dedup_preserves_unique_orders(spark):
    """
    Строки с разными order_id не должны удаляться - они уникальны.
    """
    df = spark.createDataFrame(
        [("ORD-001", "2024-01-15 10:00:00", 18.0),
         ("ORD-002", "2024-01-15 11:00:00", 50.0),
         ("ORD-003", "2024-01-15 12:00:00", 30.0)],
        schema="order_id STRING, updated_at STRING, order_total DOUBLE"
    )
    result = deduplicate_orders(df)
    assert result.count() == 3, "Unique orders should all be preserved"

# tests/schema/test_schema_contracts.py

import pytest
from pyspark.testing.utils import assertSchemaEqual
from pyspark.sql.types import (
    StructType, StructField, StringType, IntegerType,
    DoubleType, TimestampType, DateType, BooleanType
)
from src.orders_pipeline import clean_orders, calculate_order_total, enrich_with_customer

# ── Ожидаемые схемы (Data Contracts) ─────────────────────────────────────────

# Контракт на Bronze Orders схему (от источника)
BRONZE_ORDERS_SCHEMA = StructType([
    StructField("order_id",    StringType(),    nullable=True),
    StructField("customer_id", StringType(),    nullable=False),
    StructField("quantity",    IntegerType(),   nullable=True),
    StructField("unit_price",  DoubleType(),    nullable=True),
    StructField("discount",    DoubleType(),    nullable=True),
    StructField("updated_at",  StringType(),    nullable=True),
])

# Контракт на Silver Orders схему (наш output)
SILVER_ORDERS_SCHEMA = StructType([
    StructField("order_id",      StringType(),    nullable=True),
    StructField("customer_id",   StringType(),    nullable=False),
    StructField("quantity",      IntegerType(),   nullable=True),
    StructField("unit_price",    DoubleType(),    nullable=True),
    StructField("discount",      DoubleType(),    nullable=True),
    StructField("updated_at",    StringType(),    nullable=True),
    StructField("order_total",   DoubleType(),    nullable=True),
    StructField("customer_name", StringType(),    nullable=True),
    StructField("segment",       StringType(),    nullable=True),
])


def test_clean_orders_output_schema(spark):
    """
    clean_orders() не должна изменять схему (только фильтрует строки).
    Тест защищает от случайного добавления/удаления колонок при рефакторинге.
    """
    input_df = spark.createDataFrame(
        [("ORD-001", "CUST-01", 2, 10.0, 0.1, "2024-01-15")],
        schema=BRONZE_ORDERS_SCHEMA
    )
    result = clean_orders(input_df)
    assertSchemaEqual(result.schema, BRONZE_ORDERS_SCHEMA, checkNullable=False)


def test_calculate_order_total_adds_column(spark):
    """
    calculate_order_total() должна добавить order_total DOUBLE.
    Если тип изменится (напр. Decimal вместо Double) - тест упадёт.
    """
    input_df = spark.createDataFrame(
        [("ORD-001", "CUST-01", 2, 10.0, 0.1, "2024-01-15")],
        schema=BRONZE_ORDERS_SCHEMA
    )
    result = calculate_order_total(input_df)

    assert "order_total" in result.columns, "order_total column must be present"
    order_total_field = result.schema["order_total"]
    assert isinstance(order_total_field.dataType, DoubleType), \
        f"order_total must be DOUBLE, got {order_total_field.dataType}"


def test_schema_fail_on_missing_column(spark):
    """
    Тест специально ПРОВАЛИВАЕТСЯ если отсутствует обязательная колонка.

    Имитирует ситуацию schema drift:
    источник прислал данные без customer_id (переименовал в 'cid').
    """
    broken_schema_df = spark.createDataFrame(
        [("ORD-001", "CUST-01", 2, 10.0)],
        schema="order_id STRING, cid STRING, quantity INT, unit_price DOUBLE"
        # Нет 'customer_id', нет 'discount', нет 'updated_at'
    )

    # Пытаемся применить функцию к DataFrame с неправильной схемой
    with pytest.raises(Exception):
        # clean_orders ожидает 'discount' колонку для coalesce
        # На неправильной схеме должна упасть с AnalysisException
        clean_orders(broken_schema_df).collect()  # .collect() триггерит вычисление


def test_full_pipeline_output_schema(spark):
    """
    End-to-end схема-тест: после всей цепочки трансформаций
    схема должна точно соответствовать SILVER_ORDERS_SCHEMA.
    """
    orders = spark.createDataFrame(
        [("ORD-001", "CUST-01", 2, 10.0, 0.1, "2024-01-15")],
        schema=BRONZE_ORDERS_SCHEMA
    )
    customers = spark.createDataFrame(
        [("CUST-01", "Alice", "Premium")],
        schema="customer_id STRING, customer_name STRING, segment STRING"
    )

    result = enrich_with_customer(calculate_order_total(clean_orders(orders)), customers)

    assertSchemaEqual(result.schema, SILVER_ORDERS_SCHEMA, checkNullable=False)

# tests/unit/test_validation.py

import pytest
from pyspark.sql.utils import AnalysisException
from src.orders_pipeline import calculate_order_total


def test_fails_on_missing_required_column(spark):
    """
    Если входной DataFrame не содержит обязательную колонку 'quantity',
    функция должна упасть с понятной ошибкой - не с NullPointerException.

    pytest.raises(): тест ПРОХОДИТ если внутри блока произошло указанное исключение.
    Тест ПАДАЕТ если исключения не было (функция не упала - баг!).
    """
    df_missing_column = spark.createDataFrame(
        [("ORD-001", 10.0)],
        schema="order_id STRING, unit_price DOUBLE"
        # Нет quantity, discount - обязательных для расчёта
    )

    with pytest.raises(AnalysisException) as exc_info:
        calculate_order_total(df_missing_column).collect()

    # Дополнительно: проверяем, что сообщение об ошибке понятное
    assert "quantity" in str(exc_info.value).lower() or \
           "UNRESOLVED_COLUMN" in str(exc_info.value), \
           "Error message should mention the missing column"


def test_custom_validation_error(spark):
    """
    Пример кастомной валидации с понятным сообщением об ошибке.
    """
    from src.orders_pipeline import validate_and_clean

    df_empty = spark.createDataFrame([], schema="order_id STRING, quantity INT, unit_price DOUBLE, discount DOUBLE")

    with pytest.raises(ValueError, match="DataFrame is empty"):
        validate_and_clean(df_empty)  # Должна падать с ValueError, а не молча возвращать пустой DF

# tests/integration/test_pipeline_e2e.py

import pytest
import tempfile
import os
from pathlib import Path
from pyspark.testing import assertDataFrameEqual


@pytest.fixture(scope="function")
def tmp_dir():
    """Временная директория для тестовых Parquet/Delta файлов."""
    with tempfile.TemporaryDirectory() as tmpdir:
        yield tmpdir


def test_write_and_read_parquet_roundtrip(spark, tmp_dir):
    """
    Integration-тест: запись в Parquet и чтение обратно.
    Проверяет корректность сериализации схемы.

    Это простейший integration-тест - никакого реального S3/HDFS,
    только локальная файловая система.
    """
    original = spark.createDataFrame(
        [("ORD-001", 18.0), ("ORD-002", 50.0)],
        schema="order_id STRING, order_total DOUBLE"
    )

    output_path = os.path.join(tmp_dir, "orders_output")
    original.write.parquet(output_path)

    read_back = spark.read.parquet(output_path)

    assertDataFrameEqual(original, read_back, checkRowOrder=False)


def test_idempotent_partition_overwrite(spark, tmp_dir):
    """
    Integration-тест: идемпотентность OVERWRITE.
    Запустим пайплайн дважды → результат должен быть одинаковым (нет дублей).
    """
    output_path = os.path.join(tmp_dir, "silver_orders")

    orders = spark.createDataFrame(
        [("ORD-001", "2024-01-15", 18.0),
         ("ORD-002", "2024-01-15", 50.0)],
        schema="order_id STRING, order_date STRING, order_total DOUBLE"
    )

    # Первый запуск
    orders.write.partitionBy("order_date").mode("overwrite").parquet(output_path)

    # Второй запуск (повторный - должен быть идемпотентным)
    orders.write.partitionBy("order_date").mode("overwrite").parquet(output_path)

    result = spark.read.parquet(output_path)
    assert result.count() == 2, \
        "Idempotent overwrite should not create duplicates. " \
        f"Expected 2 rows, got {result.count()}"

def test_bronze_to_silver_contract(spark, tmp_dir):
    """
    Contract test: Silver читает Bronze и ожидает определённую схему.
    Если Bronze изменит схему → этот тест упадёт → инженер узнает до prod.
    """
    # Симулируем Bronze данные (как они приходят от источника)
    bronze_path = os.path.join(tmp_dir, "bronze")
    bronze_data = spark.createDataFrame(
        [("ORD-001", "CUST-01", 2, 10.0, 0.1, "2024-01-15 10:00:00"),
         ("ORD-002", "CUST-02", 1, 50.0, 0.0, "2024-01-15 11:00:00")],
        schema="order_id STRING, customer_id STRING, quantity INT, "
               "unit_price DOUBLE, discount DOUBLE, updated_at STRING"
    )
    bronze_data.write.parquet(bronze_path)

    # Silver пайплайн читает Bronze
    bronze_read = spark.read.parquet(bronze_path)

    # Применяем Silver трансформации
    from src.orders_pipeline import clean_orders, calculate_order_total
    silver_result = calculate_order_total(clean_orders(bronze_read))

    # Контракт: Silver должен содержать order_total
    assert "order_total" in silver_result.columns
    assert silver_result.filter("order_total IS NULL").count() == 0, \
        "All Silver rows should have non-NULL order_total"

# .github/workflows/test-pyspark.yml
# Запускается при каждом Push/PR в main

name: PySpark Unit Tests

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

jobs:
  test:
    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"

      - name: Set up Java 11
        # PySpark требует Java! Без этого шага - ошибка JAVA_HOME
        uses: actions/setup-java@v4
        with:
          distribution: "temurin"
          java-version: "11"

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

      - name: Install dependencies
        run: |
          python -m pip install --upgrade pip
          pip install -r requirements.txt
          pip install -r requirements-dev.txt

      - name: Run unit tests
        env:
          # Spark логи: только WARN/ERROR в CI - без INFO спама
          SPARK_LOCAL_IP: "127.0.0.1"
          PYSPARK_PYTHON: python3
        run: |
          pytest tests/unit/ \
            -v \
            --tb=short \
            --junitxml=test-results/unit-results.xml \
            --cov=src \
            --cov-report=xml:coverage.xml \
            --cov-report=term-missing

      - name: Run schema tests
        run: |
          pytest tests/schema/ -v --tb=short

      - name: Upload test results
        uses: actions/upload-artifact@v4
        if: always()
        with:
          name: test-results
          path: test-results/

      - name: Upload coverage
        uses: codecov/codecov-action@v4
        with:
          files: coverage.xml

# Проблема 1: SparkSession создаётся на каждый тест (медленно)
# Решение: scope="session" в conftest.py (уже показано выше)

# Проблема 2: Слишком много партиций (shuffle overhead на маленьких данных)
# В conftest.py уже есть:
# .config("spark.sql.shuffle.partitions", "4")  # Вместо default 200

# Проблема 3: Тесты зависят от порядка выполнения (shared state)
# Решение: scope="function" для fixture данных (уже показано)

# Проблема 4: Тяжёлые integration-тесты запускаются в CI на каждый commit
# Решение: маркировка тестов

# pytest.ini (добавить):
# markers =
#     integration: Integration tests (slow, require storage)
#     unit: Unit tests (fast, no IO)

# Пример маркировки:
import pytest

@pytest.mark.unit
def test_clean_orders_fast(spark):
    pass

@pytest.mark.integration
@pytest.mark.slow
def test_pipeline_with_real_storage(spark):
    pass

# Запуск только быстрых unit-тестов:
# pytest tests/ -m "unit and not slow"

# Запуск только integration тестов (например, на PR в main):
# pytest tests/ -m "integration"

"""
СЦЕНАРИЙ: разработчик сделал "небольшой рефакторинг" формулы расчёта скидки.
Было: total = qty * price * (1 - discount)
Стало: total = qty * price - discount  ← БАГ! discount теперь трактуется как сумма, а не процент

В production: отчёты о выручке показывают на 10% больше для Premium-клиентов.
Бизнес паникует. Разбор занимает 2 дня.

С unit-тестами: CI падает немедленно.
"""

# Версия с багом
def calculate_order_total_BUGGY(df: DataFrame) -> DataFrame:
    return df.withColumn(
        "order_total",
        F.round(
            F.col("quantity") * F.col("unit_price") - F.col("discount"),  # БАГИ!
            2
        )
    )

# Тест немедленно покажет ошибку:
def test_catches_formula_bug(spark):
    """Этот тест УПАДЁТ на buggy версии."""
    df = spark.createDataFrame(
        [("ORD-001", 2, 10.0, 0.1)],
        schema="order_id STRING, quantity INT, unit_price DOUBLE, discount DOUBLE"
    )

    # Правильный расчёт: 2 * 10 * (1 - 0.1) = 18.0
    # Багованный расчёт: 2 * 10 - 0.1 = 19.9

    result = calculate_order_total_BUGGY(df)
    total = result.select("order_total").collect()[0]["order_total"]

    assert abs(total - 18.0) < 0.01, \
        f"order_total should be 18.0 (qty*price*(1-discount)), got {total}. " \
        f"Check the discount formula: it should be multiplicative (1-discount), " \
        f"not subtractive (-discount)."
    # AssertionError: order_total should be 18.0, got 19.9.
    # Check the discount formula: it should be multiplicative...
    # → CI падает → PR не мержится → production спасён

"""
Структура полного test suite для Orders ETL Pipeline.

Файлы:
  tests/conftest.py                   - SparkSession, shared fixtures
  tests/unit/test_clean_orders.py     - 5 тестов, покрытие 95%
  tests/unit/test_calculate_total.py  - 4 теста, все граничные случаи
  tests/unit/test_enrich.py           - 4 теста, включая join ловушки
  tests/unit/test_dedup.py            - 4 теста, включая NULL в sort
  tests/schema/test_contracts.py      - 5 schema-тестов
  tests/integration/test_e2e.py       - 3 integration-теста

Итого: 25 тестов
Время выполнения: ~15-20 секунд (SparkSession ~5с + тесты ~10-15с)
"""

# Запуск всех тестов:
# pytest tests/ -v

# Только unit-тесты (быстро, для локальной разработки):
# pytest tests/unit/ tests/schema/ -v

# С coverage отчётом:
# pytest tests/unit/ --cov=src --cov-report=html

# Конкретный тест для отладки:
# pytest tests/unit/test_enrich.py::test_enrich_deduplicates_customers_before_join -v -s

"""
Чек-лист перед написанием тестов.
"""

TESTING_READINESS_CHECKLIST = {
    "pure_functions": {
        "check": "Трансформации вынесены в отдельные функции DataFrame → DataFrame",
        "verify": "Есть ли функции, которые можно вызвать без SparkSession.read/write?",
    },
    "io_separated": {
        "check": "IO (read/write) отделён от трансформационной логики",
        "verify": "Можно ли тестировать clean/transform/enrich без файлов на диске?",
    },
    "no_globals": {
        "check": "SparkSession не создаётся на уровне модуля",
        "verify": "Нет spark = SparkSession... вне функций?",
    },
    "deterministic": {
        "check": "Трансформации детерминированы (нет current_timestamp(), rand() в ключах)",
        "verify": "Один и тот же input всегда даёт один и тот же output?",
    },
    "explicit_schema": {
        "check": "Схемы колонок явные (не infer-schema при чтении)",
        "verify": "Используется StructType или явная строка схемы в createDataFrame?",
    },
    "error_messages": {
        "check": "Ошибки содержат понятное описание причины",
        "verify": "AssertionError с 'Expected X, got Y' лучше чем просто assert X == Y",
    },
    "edge_cases_covered": {
        "check": "Тесты покрывают NULL, empty DataFrame, дубли",
        "verify": "Есть ли тест с NULL в ключевых колонках?",
    },
    "assertDataFrameEqual": {
        "check": "Используется assertDataFrameEqual вместо collect()==",
        "verify": "checkRowOrder=False везде?",
    },
}