PySpark单元测试Mock DataFrame时MagicMock与int比较报错如何解决

你通过MagicMock模拟得到的spark_df本身是Mock实例，它调用count()方法后默认返回的还是MagicMock类型对象，不是整数，因此无法和int类型的0做比较，直接触发类型错误。

直接为Mock出来的DataFrame对象的count方法指定返回值即可，修改后的测试用例代码如下：

sys.modules["pyspark.sql"] = MagicMock()

def test_process_batch():
    df = (
        [
            (1, "foo"),
            (2, "bar"),
        ],
        ["id", "label"]
    )
    from pyspark.sql import SparkSession
    spark = SparkSession.builder.getOrCreate()
    spark_df = spark.createDataFrame(df)
    # 新增这行，指定count方法返回测试数据的实际行数
    spark_df.count.return_value = 2

    process_batch(spark_df, "123")
    assert True

如果需要测试count() <=0的分支逻辑，只需要把return_value的值改为0即可覆盖对应分支。

如果需要更贴近真实执行逻辑，不想完全Mock DataFrame的行为，也可以去掉对pyspark.sql的全局Mock，本地启动测试SparkSession构造真实的小DataFrame传入测试即可，不过这种方案执行速度会比纯Mock慢。

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