为数据流水线做单元测试

Python 中的 ETL 与 ELT

Jake Roach

Data Engineer

用单元测试校验数据流水线

单元测试：

常用于软件工程流程
确保代码按预期运行
帮助校验数据

典型的数据流水线校验框架，包括单元测试与端到端测试。

用于单元测试的 pytest

from pipeline import extract, transform, load

# 构建单元测试，断言 clean_stock_data 的类型
def test_transformed_data():
    raw_stock_data = extract("raw_stock_data.csv")
    clean_stock_data = transform(raw_data)
    assert isinstance(clean_stock_data, pd.DataFrame)

> python -m pytest

test_transformed_data .                                                     [100%]
================================ 1 passed in 1.17s ===============================

assert 与 isinstance

pipeline_type = "ETL"

# 检查 pipeline_type 是否是 str 的实例
isinstance(pipeline_type, str)

True

# 断言流水线的确等于 "ETL"
assert pipeline_type == "ETL"

# 结合 assert 与 isinstance
assert isinstance(pipeline_type, str)

AssertionError

pipeline_type = "ETL"

# 触发一个 AssertionError
assert isinstance(pipeline_type, float)

Traceback (most recent call last):
  File "<stdin>", line 4, in <module>
AssertionError

用 fixture 模拟流水线组件

import pytest

@pytest.fixture()
def clean_data():
    raw_stock_data = extract("raw_stock_data.csv")
    clean_stock_data = transform(raw_data)
    return clean_stock_data

def test_transformed_data(clean_data):
    assert isinstance(clean_data, pd.DataFrame)

对 DataFrame 做单元测试

def test_transformed_data(clean_data):
    # 在此加入更多断言
    ...

    # 检查列数
    assert len(clean_data.columns) == 4

    # 检查某列的下界
    assert clean_data["open"].min() >= 0

    # 用 "and" 链接，检查列值范围
    assert clean_data["open"].min() >= 0 and clean_data["open"].max() <= 1000

让我们一起练习吧！

Python 中的 ETL 与 ELT