Databricks中运行Unittest测试代码出现AttributeError报错求助

问题场景 #

在Databricks环境中运行以下PySpark Unittest代码时，出现了AttributeError: 'DummyMod' object has no attribute 'xxx'和SystemExit: True的报错：

原测试代码 #

import unittest
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, count, when

# Create the DataFrame
spark_df = spark.sql('select * from dna_scr_l3.vw_dna_cdm_md_bom')

class MdbomTest(unittest.TestCase):

    def test_blank_spaces_count(self):
        # Apply the blank spaces count operation and retrieve the result
        blank_spaces_count_df = spark_df.select([
            count(when((col(c).substr(1, 1) == ' ') & (col(c).substr(-1, 1) == ' '), c)).alias(c)
            for c in spark_df.columns
        ])
        
        # Perform assertions
        for column in blank_spaces_count_df.columns:
            self.assertEqual(blank_spaces_count_df.collect()[0][column], 0)

    def test_null_persentage_in_col(self):
        col_null_cnt_df = spark_df.select([count(when(col(c).isNull(), c)).alias(c) for c in spark_df.columns])
        
        # Perform assertions
        for column in col_null_cnt_df.columns:
            self.assertEqual(col_null_cnt_df.collect()[0][column], 0)

if __name__ == "__main__":
    unittest.main()

报错信息 #

ERROR: /databricks/kernel-connections/76ddb9359db2cb425c430527d257d2fe22bc38113daa396b5589d6712ace288c (unittest.loader._FailedTest)
AttributeError: 'DummyMod' object has no attribute '/databricks/kernel-connections/76ddb9359db2cb425c430527d257d2fe22bc38113daa396b5589d6712ace288c' #

运行1个测试，耗时0.005秒
失败（错误数=1）
/databricks/python/lib/python3.9/site-packages/IPython/core/interactiveshell.py:3386: UserWarning: 退出请使用'exit'、'quit'或Ctrl-D。
warn("To exit: use 'exit', 'quit', or Ctrl-D.", stacklevel=1)
SystemExit: True
已发生异常，使用%tb查看完整回溯信息。
SystemExit: True

解决方案 #

1. 移全局DataFrame初始化到测试类内部 #

全局作用域中直接执行spark.sql会导致Databricks模块加载时出现异常，需将DataFrame初始化逻辑放到测试类的setUpClass（类级初始化，仅执行一次）或setUp（方法级初始化，每次测试执行）中。

2. 修改unittest启动方式 #

在Databricks交互式环境中，unittest.main()会触发退出解释器的操作，需改用TestLoader加载测试用例并运行，或添加exit=False参数避免SystemExit。

3. 优化数据读取逻辑（可选） #

用first()替代collect()[0]，减少驱动端数据传输，提升效率。

修改后的完整代码 #

import unittest
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, count, when

class MdbomTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # 仅在测试类初始化时加载一次DataFrame
        cls.spark_df = spark.sql('select * from dna_scr_l3.vw_dna_cdm_md_bom')

    def test_blank_spaces_count(self):
        blank_spaces_count_df = self.spark_df.select([
            count(when((col(c).substr(1, 1) == ' ') & (col(c).substr(-1, 1) == ' '), c)).alias(c)
            for c in self.spark_df.columns
        ])
        
        result = blank_spaces_count_df.first()
        for column in blank_spaces_count_df.columns:
            self.assertEqual(result[column], 0)

    def test_null_persentage_in_col(self):
        col_null_cnt_df = self.spark_df.select([count(when(col(c).isNull(), c)).alias(c) for c in self.spark_df.columns])
        
        result = col_null_cnt_df.first()
        for column in col_null_cnt_df.columns:
            self.assertEqual(result[column], 0)

if __name__ == "__main__":
    # 方案1：使用TestLoader和TextTestRunner运行
    runner = unittest.TextTestRunner(verbosity=2)
    test_suite = unittest.TestLoader().loadTestsFromTestCase(MdbomTest)
    runner.run(test_suite)
    
    # 方案2：使用unittest.main并禁用退出
    # unittest.main(argv=['ignore-argv'], exit=False)

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