PySpark:模拟udf函数

Question

我正在尝试测试我编写的一个功能，它的一部分是使用一个调用外部库的udf，我想模拟对该库的调用，并返回化装值，这可能吗？

我想测试的代码如下所示：

import some_library

def my_foo(df):
    my_udf = udf(udf_foo)
    return df.withColumn("new_col", my_udf)

def udf_foo(x):
    return some_library.foo(x)

我想模拟some_library.foo，并相应地返回值。我的测试看起来像这样：

@patch('some_libaray.foo')
def test(some_libaray_mock):
   some_libaray_mock.side_effect = mocking_some_libaray_foo
   # rest of test

def mocking_some_libaray_foo(x):
   if x == 1:
     return 2
   ...

Answer 1

我不能解决这个具体的问题，但我确实找到了解决这个问题的办法。

在另一个类/服务中扭曲我的udf函数，并将它传递给我的测试类，这样它就会从它创建udf，现在模拟的是easy.

• Extracting类之外的函数，并从我的测试类覆盖它，这不是一个好的解决方案，但它的工作没有很多变化。