import pytest
from pyspark.sql import functions as F
from my_transforms import categorize_age

# Each tuple is (input_age, expected_category)
@pytest.mark.parametrize("age, expected", [
    (10, "child"),
    (17, "child"),
    (18, "adult"),
    (64, "adult"),
    (65, "senior"),
    (90, "senior"),
])
def test_categorize_age(spark, age, expected):
    # Create a one-row DataFrame with this age
    df = spark.createDataFrame([(1, age)], ["id", "age"])
    
    result = categorize_age(df)  # adds "age_category" column
    category = result.collect()[0]["age_category"]
    
    assert category == expected, f"Age {age}: expected '{expected}', got '{category}'"

# Pytest runs this test 6 times automatically — once per parameter set
# Output shows: test_categorize_age[10-child] PASSED
#               test_categorize_age[65-senior] PASSED   etc.

[pytest]
testpaths = tests
python_files = test_*.py
python_functions = test_*

# Custom markers to tag test types
markers =
    unit: fast isolated unit tests
    integration: tests that use real Spark + files
    slow: tests that take a long time

# Run only unit tests: pytest -m unit
# Run all tests:       pytest
# Run with verbose:    pytest -v

import pytest

@pytest.mark.unit
def test_remove_nulls(spark):
    df = spark.createDataFrame([(1, None), (2, "Alice")], ["id", "name"])
    result = df.dropna()
    assert result.count() == 1

@pytest.mark.integration
def test_read_parquet_and_transform(spark, tmp_path):
    # This test reads a real parquet file
    ...

import os
from pathlib import Path

# Pattern 1: Inline test data (best for unit tests)
def test_transform_inline(spark):
    df = spark.createDataFrame([
        (1, "apple", 100),
        (2, "banana", 200),
    ], ["id", "product", "revenue"])
    
    result = df.filter("revenue > 150")
    assert result.count() == 1


# Pattern 2: Test files from disk (best for integration tests)
TEST_DATA_DIR = Path(__file__).parent / "data"

def test_read_csv(spark):
    df = spark.read.csv(str(TEST_DATA_DIR / "orders.csv"), header=True, inferSchema=True)
    assert df.count() > 0
    assert "order_id" in df.columns


# Pattern 3: Temp directory (pytest's tmp_path fixture)
def test_write_and_read_parquet(spark, tmp_path):
    df = spark.range(10)  # creates DataFrame with id column 0-9
    
    out = str(tmp_path / "output")
    df.write.parquet(out)
    
    result = spark.read.parquet(out)
    assert result.count() == 10
    # tmp_path is automatically cleaned up by pytest after the test

# conftest.py
from pyspark.sql.types import (StructType, StructField, IntegerType,
                                   StringType, DoubleType, DateType)
from datetime import date

@pytest.fixture
def orders_df(spark):
    """Sample orders DataFrame with explicit schema."""
    schema = StructType([
        StructField("order_id",   IntegerType(), nullable=False),
        StructField("customer_id", IntegerType(), nullable=True),
        StructField("amount",      DoubleType(),  nullable=True),
        StructField("status",      StringType(),  nullable=True),
        StructField("order_date",  DateType(),    nullable=True),
    ])
    data = [
        (1001, 1, 250.0,  "COMPLETED", date(2024, 1, 15)),
        (1002, 2, 89.99, "PENDING",   date(2024, 1, 16)),
        (1003, 1, 500.0,  "COMPLETED", date(2024, 1, 17)),
        (1004, 3, None,   "FAILED",    date(2024, 1, 18)),
    ]
    return spark.createDataFrame(data, schema)


@pytest.fixture
def products_df(spark):
    data = [(1, "Laptop", 999.99), (2, "Mouse", 29.99)]
    return spark.createDataFrame(data, ["product_id", "name", "price"])


# A test using multiple fixtures:
def test_orders_are_all_positive(orders_df):
    # Test that all non-null amounts are positive
    negatives = orders_df.filter("amount <= 0")
    assert negatives.count() == 0

def test_completed_orders_count(orders_df):
    completed = orders_df.filter("status = 'COMPLETED'")
    assert completed.count() == 2

def test_no_rows_lost_in_join(spark):
    customers = spark.createDataFrame(
        [(1, "Alice"), (2, "Bob"), (3, "Carol")],
        ["cust_id", "name"]
    )
    orders = spark.createDataFrame(
        [(1, 100.0), (1, 200.0), (2, 150.0)],
        ["cust_id", "amount"]
    )
    
    # Left join — Carol has no orders but should still appear
    result = customers.join(orders, on="cust_id", how="left")
    
    # 3 customers × up to 2 orders each = 4 rows (Carol gets 1 null row)
    assert result.count() == 4
    
    # Verify Carol appears with null amount
    carol_rows = result.filter("name = 'Carol'")
    assert carol_rows.count() == 1
    assert carol_rows.collect()[0]["amount"] is None


def test_dedup_reduces_rows(spark):
    df = spark.createDataFrame(
        [(1, "A"), (1, "A"), (2, "B")],
        ["id", "val"]
    )
    result = df.dropDuplicates()
    assert result.count() == 2  # was 3, should be 2 after dedup


def test_filter_result_count(orders_df):
    completed = orders_df.filter("status = 'COMPLETED'")
    
    # Exact count
    assert completed.count() == 2
    
    # Range check (useful when data size isn't fully known)
    total = orders_df.count()
    assert 0 < completed.count() < total  # some, but not all

from pyspark.sql import functions as F

def assert_df_equals(df1, df2, sort_by):
    """
    Compare two DataFrames for equality.
    - Same schema
    - Same row count
    - Same data (order-independent via sort)
    """
    # Check schemas match
    assert df1.schema == df2.schema, f"Schema mismatch:\n{df1.schema}\nvs\n{df2.schema}"

    # Check row counts match
    count1, count2 = df1.count(), df2.count()
    assert count1 == count2, f"Row count mismatch: {count1} vs {count2}"

    # Sort both and compare row by row
    rows1 = df1.orderBy(sort_by).collect()
    rows2 = df2.orderBy(sort_by).collect()
    assert rows1 == rows2, "DataFrame contents do not match!"


def test_aggregate_result(spark):
    orders = spark.createDataFrame([
        (1, "A", 100.0), (2, "A", 200.0),
        (3, "B", 150.0),
    ], ["order_id", "category", "amount"])

    result = orders.groupBy("category").agg(
        F.sum("amount").alias("total_amount"),
        F.count("*").alias("order_count")
    )

    expected = spark.createDataFrame([
        ("A", 300.0, 2),
        ("B", 150.0, 1),
    ], ["category", "total_amount", "order_count"])

    assert_df_equals(result, expected, sort_by="category")

from pyspark.sql import functions as F

def test_column_values(spark):
    df = spark.createDataFrame([
        (1, "COMPLETED", 100.0),
        (2, "PENDING",   50.0),
        (3, "FAILED",    0.0),
    ], ["id", "status", "amount"])

    # 1. Check allowed values (enum check)
    allowed_statuses = {"COMPLETED", "PENDING", "FAILED", "CANCELLED"}
    actual_statuses = {row["status"] for row in df.select("status").distinct().collect()}
    assert actual_statuses.issubset(allowed_statuses), f"Invalid statuses: {actual_statuses - allowed_statuses}"

    # 2. Check no nulls in critical column
    null_count = df.filter(F.col("id").isNull()).count()
    assert null_count == 0, "id column must not have nulls"

    # 3. Check value range
    min_amount, max_amount = df.agg(
        F.min("amount"), F.max("amount")
    ).collect()[0]
    assert min_amount >= 0, "Amount cannot be negative"
    assert max_amount < 1_000_000, "Amount seems unreasonably large"

    # 4. Uniqueness check
    total_rows = df.count()
    distinct_ids = df.select("id").distinct().count()
    assert total_rows == distinct_ids, "id column must be unique"

    # 5. Non-empty string check
    empty_status = df.filter(F.trim(F.col("status")) == "").count()
    assert empty_status == 0, "status must not be empty string"

from chispa.dataframe_comparer import assert_approx_df_equality

def test_average_calculation(spark):
    # Compute an average — floating point result
    df = spark.createDataFrame(
        [(1, 100.0), (2, 200.0), (3, 300.0)],
        ["id", "revenue"]
    )
    
    from pyspark.sql import functions as F
    result = df.agg(F.avg("revenue").alias("avg_revenue"))
    
    expected = spark.createDataFrame(
        [(200.00000001,)],  # slightly off due to floating point
        ["avg_revenue"]
    )
    
    # precision=1e-3 means values within 0.001 of each other pass
    assert_approx_df_equality(result, expected, precision=1e-3)


def test_percentage_columns(spark):
    """Good for testing percentage/ratio columns."""
    result = compute_market_share(spark)  # returns % values
    expected = spark.createDataFrame(
        [("Product A", 33.33), ("Product B", 66.67)],
        ["product", "share_pct"]
    )
    
    assert_approx_df_equality(result, expected,
                              precision=0.01,          # within 0.01%
                              ignore_row_order=True)

from chispa.schema_comparer import assert_schema_equality
from pyspark.sql.types import *

def test_after_cast(spark):
    df = spark.createDataFrame(
        [("1", "100.5")],
        ["user_id", "amount"]
    )
    
    # Cast string columns to correct types
    result = df.withColumn("user_id", df["user_id"].cast(IntegerType())) \
               .withColumn("amount",  df["amount"].cast(DoubleType()))
    
    expected_schema = StructType([
        StructField("user_id", IntegerType(), nullable=True),
        StructField("amount",  DoubleType(),  nullable=True),
    ])
    
    assert_schema_equality(result.schema, expected_schema)


# Ignore nullable flags if they don't matter for your test:
from chispa.schema_comparer import assert_schema_equality_ignore_nullable

def test_schema_ignoring_nullable(spark, result_df):
    expected_schema = StructType([
        StructField("id",   IntegerType(), nullable=False),
        StructField("name", StringType(),  nullable=False),
    ])
    # Won't fail if actual has nullable=True
    assert_schema_equality_ignore_nullable(result_df.schema, expected_schema)

import pytest
from testcontainers.postgres import PostgresContainer
import psycopg2

@pytest.fixture(scope="session")
def postgres_container():
    """Start a real PostgreSQL container for the test session."""
    with PostgresContainer("postgres:15") as postgres:
        yield postgres  # container is running during tests
    # Container auto-destroyed after yield block exits


@pytest.fixture(scope="session")
def seeded_postgres(postgres_container):
    """Create tables and insert test data into the container."""
    conn = psycopg2.connect(postgres_container.get_connection_url())
    cursor = conn.cursor()
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS orders (
            order_id   INT PRIMARY KEY,
            customer   VARCHAR(100),
            amount     NUMERIC(10, 2),
            status     VARCHAR(20)
        )
    """)
    cursor.execute("""
        INSERT INTO orders VALUES
        (1, 'Alice', 250.00, 'COMPLETED'),
        (2, 'Bob',   89.99,  'PENDING'),
        (3, 'Carol', 500.00, 'COMPLETED')
    """)
    conn.commit()
    cursor.close()
    conn.close()
    return postgres_container


def test_spark_reads_postgres(spark, seeded_postgres):
    # Build JDBC URL from the container
    jdbc_url = seeded_postgres.get_connection_url().replace(
        "postgresql", "jdbc:postgresql"
    )
    
    # Read the table using Spark JDBC
    df = spark.read.format("jdbc") \
        .option("url", jdbc_url) \
        .option("dbtable", "orders") \
        .option("user", seeded_postgres.POSTGRES_USER) \
        .option("password", seeded_postgres.POSTGRES_PASSWORD) \
        .option("driver", "org.postgresql.Driver") \
        .load()
    
    # Verify data
    assert df.count() == 3
    completed = df.filter("status = 'COMPLETED'")
    assert completed.count() == 2
    
    total = df.agg({"amount": "sum"}).collect()[0][0]
    assert abs(total - 839.99) < 0.01

import pytest, json
from testcontainers.kafka import KafkaContainer
from kafka import KafkaProducer

@pytest.fixture(scope="session")
def kafka_container():
    with KafkaContainer() as kafka:
        yield kafka


def test_streaming_from_kafka(spark, kafka_container, tmp_path):
    bootstrap = kafka_container.get_bootstrap_server()
    topic = "test-orders"

    # Produce 3 test messages to Kafka
    producer = KafkaProducer(
        bootstrap_servers=bootstrap,
        value_serializer=lambda v: json.dumps(v).encode()
    )
    for i in range(3):
        producer.send(topic, {"order_id": i, "amount": i * 100})
    producer.flush()

    # Read from Kafka with Spark Streaming (batch mode)
    df = spark.read \
        .format("kafka") \
        .option("kafka.bootstrap.servers", bootstrap) \
        .option("subscribe", topic) \
        .option("startingOffsets", "earliest") \
        .load()
    
    # Parse JSON values
    from pyspark.sql.types import StructType, StructField, IntegerType
    from pyspark.sql import functions as F
    
    schema = StructType([
        StructField("order_id", IntegerType()),
        StructField("amount", IntegerType()),
    ])
    parsed = df.select(F.from_json(F.col("value").cast("string"), schema).alias("data")) \
               .select("data.*")
    
    assert parsed.count() == 3
    total = parsed.agg({"amount": "sum"}).collect()[0][0]
    assert total == 300  # 0 + 100 + 200

# pip install pydeequ
from pydeequ.checks import Check, CheckLevel
from pydeequ.verification import VerificationSuite, VerificationResult

def test_with_deequ(spark):
    df = spark.createDataFrame([
        (1, "Alice", 28,  "alice@email.com"),
        (2, "Bob",   35,  "bob@email.com"),
        (3, None,   22,  "carol@email.com"),
    ], ["id", "name", "age", "email"])

    check = (Check(spark, CheckLevel.Error, "Order Quality")
        .isComplete("id")                      # id has no nulls
        .isUnique("id")                         # id is unique
        .isNonNegative("age")                   # age >= 0
        .isContainedIn("name", ["Alice", "Bob", "Carol"])  # allowed names
        .hasCompleteness("name", lambda p: p >= 0.8)  # 80%+ non-null
    )

    result = VerificationSuite(spark).onData(df).addChecks([check]).run()
    df_result = VerificationResult.checkResultsAsDataFrame(spark, result)
    
    # Show detailed results
    df_result.show(truncate=False)
    
    # Fail test if any constraint failed
    failed = df_result.filter("constraint_status = 'Failure'")
    assert failed.count() == 0, f"DQ constraints failed!"

# checks.yml
checks for orders:
  # Row count must be between 100 and 1 million
  - row_count:
      between [100, 1000000]

  # order_id must have zero nulls
  - missing_count(order_id) = 0

  # status only allows these values
  - invalid_count(status):
      valid values: [COMPLETED, PENDING, FAILED, CANCELLED]
      must be 0

  # amount must be positive
  - min(amount) > 0

  # customer_email must match regex
  - invalid_count(customer_email):
      valid regex: '^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$'
      must be 0

  # Duplicate check
  - duplicate_count(order_id) = 0

# pip install soda-core-spark-df
from soda.scan import Scan

def test_soda_quality(spark):
    df = spark.read.parquet("data/orders.parquet")
    
    scan = Scan()
    scan.set_scan_definition_name("orders-test")
    scan.set_data_source_name("spark_df")
    scan.add_spark_session(spark)
    
    # Register the DataFrame under the name "orders"
    scan.add_dataframe_table(df, "orders")
    
    # Add the YAML check file
    scan.add_sodacl_yaml_file("checks/checks.yml")
    
    # Run checks
    scan.execute()
    
    # Assert no failures
    assert scan.get_checks_fail_count() == 0, \
        f"Soda checks failed:\n{scan.get_scan_results()}"

from dataclasses import dataclass
from typing import List, Callable
from pyspark.sql import DataFrame
from pyspark.sql import functions as F

@dataclass
class ValidationRule:
    name: str
    check: Callable[[DataFrame], bool]
    description: str


class DataValidator:
    def __init__(self, df: DataFrame):
        self.df = df
        self.rules: List[ValidationRule] = []
    
    def add_rule(self, name, check, description=""):
        self.rules.append(ValidationRule(name, check, description))
        return self  # allow chaining
    
    def validate(self):
        results = []
        for rule in self.rules:
            try:
                passed = rule.check(self.df)
                results.append({"rule": rule.name, "passed": passed,
                                 "desc": rule.description, "error": None})
            except Exception as e:
                results.append({"rule": rule.name, "passed": False,
                                 "desc": rule.description, "error": str(e)})
        return results


# Using the custom validator:
def test_custom_validator(spark, orders_df):
    results = (
        DataValidator(orders_df)
        .add_rule(
            "no_null_order_ids",
            lambda df: df.filter(F.col("order_id").isNull()).count() == 0,
            "Order IDs must not be null"
        )
        .add_rule(
            "positive_amounts",
            lambda df: df.filter("amount <= 0").count() == 0,
            "All amounts must be positive"
        )
        .add_rule(
            "valid_statuses",
            lambda df: df.filter(
                ~F.col("status").isin(["COMPLETED", "PENDING", "FAILED"])
            ).count() == 0,
            "Status must be one of allowed values"
        )
        .validate()
    )

    failures = [r for r in results if not r["passed"]]
    assert not failures, f"Validation failed: {failures}"