Databricks 中的单元测试框架

Question

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()

想要运行单元测试但出现错误

错误：/databricks/kernel-connections/76ddb9359db2cb425c430527d257d2fe22bc38113daa396b5589d6712ace288c（unittest.loader._FailedTest）

AttributeError：“DummyMod”对象没有属性“/databricks/kernel-connections/76ddb9359db2cb425c430527d257d2fe22bc38113daa396b5589d6712ace288c”

---

在 0.005 秒内运行 1 次测试

失败（错误=1） /databricks/python/lib/python3.9/site-packages/IPython/core/interactiveshell.py:3386：UserWarning：要退出：使用“exit”、“quit”或 Ctrl-D。 warn("要退出：使用 'exit'、'quit' 或 Ctrl-D。", stacklevel=1) 系统退出：真 发生异常，使用 %tb 查看完整的回溯。

系统退出：真

Answer 1

问题是，要在 Databricks 中使用

unittest.main

函数，您需要在参数列表中设置

exit = False

，否则它将终止当前的 Python shell。将

argv

作为单元素列表显式传递也是有意义的，以避免使用 Databricks 上包含用于启动 Python 子进程的参数的 s

ys.argv

。 （参见 unittest.main 的文档）：

test_runner = unittest.main(argv=[''], exit=False)
test_runner.result.printErrors()

if not test_runner.result.wasSuccessful():
  raise Exception(
    f"{len(test_runner.result.failures)} of {test_runner.result.testsRun} tests failed.")

此外，在笔记本中，

__name__

始终设置为

__main__

。

但请考虑到，默认情况下，

unittest

自动发现可能不起作用，因此您可以显式注册类/测试用例，或者编写一些代码来自动发现测试用例。

Answer 2

而不是

unittest.main()

用这个：

unittest.main(argv=[''], exit=False)