단언으로 Pipeline 단계 test하기

import pandas as pd
import numpy as np

# A transform that looks correct but has a bug:
def compute_revenue_buggy(df):
    # BUG: unit_price should be multiplied, not added
    df['revenue'] = df['quantity'] + df['unit_price']
    return df

# Without assertions, this runs silently with wrong numbers.

def drop_nulls_guarded(df, required_cols, max_drop_fraction=0.3):
    before = len(df)
    result = df.dropna(subset=required_cols)
    after = len(result)
    drop_fraction = (before - after) / before
    assert drop_fraction <= max_drop_fraction, \
        f'dropna removed {drop_fraction:.1%} of rows (limit {max_drop_fraction:.1%})'
    return result

def compute_revenue(df):
    df = df.assign(revenue=lambda d: d['quantity'] * d['unit_price'])

    # Guard: check that the new column exists and is non-null
    assert 'revenue' in df.columns, 'revenue column not created'
    assert df['revenue'].notna().all(), 'revenue has unexpected NaNs'
    return df

def validate_computed_columns(df):
    assert (df['revenue'] >= 0).all(), \
        f'Negative revenue: {df[df["revenue"] < 0]["revenue"].head().tolist()}'
    assert (df['quantity'] > 0).all(), \
        'Non-positive quantity found'
    assert df['revenue'].between(0, 1_000_000).all(), \
        'Revenue out of plausible range'
    print('Value range checks passed.')

def safe_merge(left, right, on, how='left'):
    before = len(left)
    result = left.merge(right, on=on, how=how)
    after = len(result)

    if how == 'left':
        assert after == before, \
            f'Left join increased rows from {before} to {after} — check key uniqueness in right df'
    return result

NOT_NULL_AFTER_TRANSFORM = ['order_id', 'revenue', 'category']

def post_transform_checks(df):
    for col in NOT_NULL_AFTER_TRANSFORM:
        null_count = df[col].isna().sum()
        assert null_count == 0, \
            f'{col}: {null_count} unexpected NaN values after transform'
    print('Null checks passed.')
    return df

def test_compute_revenue():
    test_df = pd.DataFrame({
        'quantity': [2, 3],
        'unit_price': [10.0, 5.0]
    })
    result = compute_revenue(test_df)

    assert 'revenue' in result.columns
    assert result['revenue'].tolist() == [20.0, 15.0]
    assert result['revenue'].dtype == float
    print('test_compute_revenue PASSED')

test_compute_revenue()

def test_empty_df():
    empty = pd.DataFrame({'quantity': [], 'unit_price': []})
    result = compute_revenue(empty)
    assert len(result) == 0, 'Empty input should produce empty output'
    assert 'revenue' in result.columns, 'Revenue column should still be created'
    print('test_empty_df PASSED')

def test_single_row():
    single = pd.DataFrame({'quantity': [1], 'unit_price': [99.0]})
    result = compute_revenue(single)
    assert result['revenue'].iloc[0] == 99.0
    print('test_single_row PASSED')

test_empty_df()
test_single_row()

def integration_test(config):
    raw = extract(config)
    clean = transform(raw, config)

    assert set(config['required_cols']).issubset(set(clean.columns))
    assert clean['order_id'].is_unique
    assert (clean['revenue'] >= 0).all()
    assert len(clean) > 0

    print(f'Integration test PASSED. Output: {clean.shape}')

integration_test(CONFIG)

# tests/test_transform.py — example structure
# import pytest, pandas as pd
# from pipeline.transform import compute_revenue, drop_nulls

# @pytest.fixture
# def sample_df():
#     return pd.DataFrame({'quantity': [2, 3], 'unit_price': [10.0, 5.0]})

# def test_revenue(sample_df):
#     result = compute_revenue(sample_df)
#     assert result['revenue'].tolist() == [20.0, 15.0]

print('Run: pytest tests/ -v  to execute all pipeline tests')

# These assertions document business rules as code:
assert (df['order_date'] >= pd.Timestamp('2020-01-01')).all(), \
    'Company was founded 2020-01-01; no orders can predate this'
assert df['region'].isin({'North', 'South', 'East', 'West', 'Central'}).all(), \
    'Only 5 sales regions are valid; new regions require config update'
print('Business rules verified as assertions.')