如何在ADF中实现Databricks Notebook函数与PySpark DF的单元测试

1. 先搞定函数的导入 #

要在测试Notebook里调用目标函数，有两种靠谱方式：

• 快速导入：用%run加载函数Notebook
  在测试Notebook最开头加一行：

%run /Workspace/路径/到/你的/函数定义Notebook

这样就能直接用目标Notebook里的所有函数，适合快速测试。

• 长期维护：封装成Databricks库
  如果函数多且要频繁复用，把函数打包成Python wheel，上传到Databricks工作区作为库，然后在测试Notebook里安装并导入：

# 安装库（如果没预装的话）
dbutils.library.install("/Workspace/路径/到/你的库文件.whl")
# 导入模块
import your_function_module

2. 函数单元测试：用unittest/pytest写用例 #

直接用Python自带的unittest框架就行，Databricks Notebook支持直接跑测试。举个例子：
假设你有个处理字符串的函数clean_string(s)，测试代码如下：

import unittest

class TestDataFunctions(unittest.TestCase):
    def test_clean_string(self):
        # 测试正常字符串
        self.assertEqual(clean_string("  Hello World!  "), "Hello World!")
        # 测试空字符串
        self.assertEqual(clean_string(""), "")
        # 测试特殊字符
        self.assertEqual(clean_string("a@#$b"), "ab")

# 运行测试，exit=False避免Notebook直接退出
if __name__ == '__main__':
    unittest.main(argv=[''], exit=False)

运行后会输出测试通过率、失败用例明细，一目了然。

3. 源文件与DataFrame计数匹配测试 #

核心是分别拿到源文件的有效记录数和处理后DataFrame的行数，然后对比：

3.1 统计源文件记录数 #

根据源文件类型选方法：

• CSV/文本文件（带表头）：

source_path = "abfss://容器名@存储账户.dfs.core.windows.net/源文件路径.csv"
# 读取所有行后减1（跳过表头）
source_record_count = spark.read.text(source_path).count() - 1

• Parquet/ORC文件：直接用spark读取后count（这类文件自带元数据，计数很快）：

source_record_count = spark.read.parquet(source_path).count()

3.2 统计DataFrame记录数 #

# df是你处理后的目标DataFrame
df_record_count = df.count()

3.3 断言匹配 #

用断言抛出异常，ADF会识别为任务失败：

assert source_record_count == df_record_count, \
    f"计数不匹配：源文件{source_record_count}条，DataFrame{df_record_count}条"

如果计数对不上，Notebook会直接报错，ADF管道会终止并标记失败。

4. 在ADF里调度测试Notebook #

• 打开ADF，新建Notebook活动，选择你的Databricks工作区、集群，以及测试Notebook的路径。
• 如果需要动态传参（比如不同环境的源文件路径），可以在ADF活动里设置参数，在Notebook里用dbutils.widgets.get("参数名")获取。
• 把这个测试活动放到数据处理任务的后面，作为数据质量校验环节。如果测试失败，整个管道停掉，避免脏数据流入下游。
• 配置告警：在ADF里设置警报规则，当测试活动失败时，自动发邮件通知运维或开发人员。

5. 额外优化：记录测试结果 #

可以把测试结果写入到日志表或ADLS文件里，方便回溯：

# 假设用Delta表存测试日志
from pyspark.sql import Row
import datetime

test_result = Row(
    test_name="函数与计数测试",
    status="PASS" if source_record_count == df_record_count else "FAIL",
    message=f"源文件{source_record_count}条，DataFrame{df_record_count}条",
    test_time=datetime.datetime.now()
)

# 写入Delta表
spark.createDataFrame([test_result]).write.mode("append").saveAsTable("test_logs.test_results")

内容的提问来源于stack exchange，提问作者harshith