PySpark: Mocking a udf function

user_9806663 0
Viewed 223

I'm trying to test a functionality I written, part of it is to use a udf that calls an outside library, I want to mock the call to that library, and return costume values, is it possible?

my code I want to test looks something like that:

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)

I want to mock some_library.foo, and return values accordingly. My test look something like this:

@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
   ...
pyspark apache-spark python-unittest mocking unit-testing
1 Answers

I couldn't solve this specific problem, but I did mange to find solution to that problem.

  • Warp my udf function inside another class/service and pass it to my tested class so it will create udf from it, and now mocking is easy.
  • Extracting the function outside of the class and overwrite it from my test class, it isn't a good solution but it works without a lot of changes.
user_98066630
Related