使用Pytest测试Spark代码时，MagicMock与int比较报错问题

问题详情 #

在编写Spark数据处理代码的单元测试时，被测代码包含f.size(f.collect_set(col)) > 1的比较逻辑，执行测试时触发TypeError，错误提示：

E TypeError: '>' not supported between instances of 'MagicMock' and int

被测代码 #

from pyspark.sql import DataFrame
import pyspark.sql.functions as f

def check(
    df: DataFrame,
) -> DataFrame:
    check_cols = ["a","b"]
    cols = [
        f.when(f.size(f.collect_set(col)) > 1, f.lit(None))
        .otherwise(f.first(col))
        .alias(col)
        for col in check_cols
    ]
    new_cols = df.groupby("col1").agg(
        f.to_json(f.struct(*cols)).alias(
            "new_col"
        )
    )
    
    return new_cols

测试代码 #

from unittest.mock import patch, MagicMock
from pyspark.sql import DataFrame

@patch("pyspark.sql.functions.to_json")
@patch("pyspark.sql.functions.struct")
@patch("pyspark.sql.functions.when")
@patch("pyspark.sql.functions.size")
@patch("pyspark.sql.functions.collect_set")
@patch("pyspark.sql.functions.lit")
@patch("pyspark.sql.functions.first")
def test_get_metadata_columns(
    first_mock,
    lit_mock,
    collect_set_mock,
    size_mock,
    when_mock,
    struct_mock,
    to_json_mock,
):
    df = MagicMock()
    check(df)
    first_mock.assert_called_once()
    collect_set_mock.assert_called_once()
    lit_mock.assert_called_once()
    struct_mock.assert_called_once()
    to_json_mock.assert_called_once()
    size_mock.assert_called_once()
    when_mock.assert_called_once()

问题原因 #

1. size_mock返回的是MagicMock实例，代码中直接将其与整数1做>比较，但MagicMock默认未实现__gt__（大于比较）魔法方法，因此触发类型错误。
2. 被测代码中check_cols包含两个元素，循环会执行两次，first_mock、collect_set_mock等方法会被调用两次，原测试中的assert_called_once()断言会失效。

解决方案 #

步骤1：配置size_mock支持比较操作 #

创建一个带__gt__方法的Mock对象，让size_mock返回该对象，确保size(...) > 1操作能生成可被f.when()接受的Mock实例。

步骤2：修正断言调用次数 #

将被循环调用的Mock断言从assert_called_once()改为assert_called_times(2)，单次调用的Mock保持原断言。

修改后的测试代码 #

from unittest.mock import patch, MagicMock
from pyspark.sql import DataFrame

@patch("pyspark.sql.functions.to_json")
@patch("pyspark.sql.functions.struct")
@patch("pyspark.sql.functions.when")
@patch("pyspark.sql.functions.size")
@patch("pyspark.sql.functions.collect_set")
@patch("pyspark.sql.functions.lit")
@patch("pyspark.sql.functions.first")
def test_get_metadata_columns(
    first_mock,
    lit_mock,
    collect_set_mock,
    size_mock,
    when_mock,
    struct_mock,
    to_json_mock,
):
    # 配置size返回值支持大于比较
    size_col_mock = MagicMock()
    size_col_mock.__gt__.return_value = MagicMock()
    size_mock.return_value = size_col_mock
    
    df = MagicMock()
    check(df)
    
    # 循环执行两次，对应Mock调用次数为2
    first_mock.assert_called_times(2)
    collect_set_mock.assert_called_times(2)
    size_mock.assert_called_times(2)
    when_mock.assert_called_times(2)
    
    # 仅被调用一次的Mock断言
    lit_mock.assert_called_once()
    struct_mock.assert_called_once()
    to_json_mock.assert_called_once()

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