pytest to compare mocked spark.sql.functions with int using >

Question

I am trying to write unit test where in the code I am comparing values. The test fails with not able to compare MagicMock with int. Below is the code under test:

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

The test case I'm writing is the following:

@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(
    to_json_mock,
    struct_mock,
    when_mock,
    size_mock,
    collect_set_mock,
    lit_mock,
    first_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()

I get the error:

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