Stop mocking me! Unit tests in PySpark using Python’s mock library
By: Serena Mcdonnell
Tags: Data Science, Pyspark, Unit Testing, Mock

Testing.

Fundamental in software development, and often overlooked by data scientists, but important. In this post, I’ll show how to do unit testing in PySpark using Python’s unittest.mock library. I’ll do this from a data scientist’s perspective- to me that means that I won’t go into the software engineering details. I present just what you need to know.

First, a (semi) relevant clip from Family Guy:

What is a unit test? What’s a mock?

A unit test is a way to test pieces of code to make sure things work as they should. The unittest.mock library in Python allows you to replace parts of your code with mock objects and make assertions about how they’ve been used. A “mock” is an object that does as the name says- it mocks the attributes of the objects/variables in your code.

The end goal: testing spark.sql(query)

A simple way to create a dataframe in PySpark is to do the following:

df = spark.sql("SELECT * FROM table")

Although it’s simple, it should be tested.

The code and problem set up

Let’s say we work for an e-commerce clothing company, and our goal is to create a product similarity table that has been filtered on some conditions, and write it to HDFS.

Assume we have the following tables:

1. Products. Columns: “item_id”, “category_id”. 2. Product_similarity (unfiltered). Columns: “item_id_1”, “item_id_2”, “similarity_score”.

(Assume the score in product_similarity is between 0 and 1, where the closer the score is to 1, the more similar the items. See my post on similarity metrics for more details if you’re interested).

Looking at pairs of items and their scores is as simple as:

The where clause is to remove rows comparing an item to itself. The score will always be 1. How boring!

But what if we want to create a table that shows us similarity of items that are in the same category? What if we don’t care to compare shoes to scarves, but we want to compare shoes to shoes and scarves to scarves? This is a bit more complicated, and requires us to join the “products” and “product_similarity” tables.

The query would then become:

We may also want to get the most N similar items for each product, so in that case, our query would become:

(Assuming we use N = 10).

Now, what if we want the option to either compare products across categories, or only within categories? We could achieve this by using a boolean variable called same_category that results in a string same_category_q that can be passed into the overall query (using .format()), and will be equal to the inner join above if our boolean same_category is True, and empty if it’s False. The query would then look like:

Let’s make it a bit more clear, and wrap this logic in a function that returns same_category_q:

So far, so good. We output the query same_category_q so we can test our function to make sure it returns what we want it to return.

Keeping in mind our final goal, we want to write a dataframe to HDFS. We can do this with the following function:

Adding the first part of the query and a main method to complete our script, we get:

The idea here is that we want to create a function, generically named test_name_of_function(), for each function in our script. We want to test that the function behaves as it should, and we ensure this by using assert statements all over the place.

Test_make_query, true and false

First, let’s test the make_query function. Recall that make_query takes in two inputs: a boolean variable and some table paths. It’ll return different values for same_category_q based on the boolean same_category. What we’re sort of doing here is like a set of if-then statements:

1. If same_category is True, then same_category_q = “INNER JOIN …” 2. If same_category is False, then same_category_q = “” (empty)

What we do is mock make_query‘s parameters and pass them in, then test that we got our desired outputs. Since test_paths is a dictionary, we don’t need to mock it. The test script is below, with comments added for extra clarity.

And it’s as simple as that!

Testing our table creation

Next up, we need to test that our create_new_table function behaves as it should. Stepping through the function, we see that it does several things, with several opportunities to make some assertions and mocks. Note that whenever we have something like df.write.save.something.anotherthing, we need to mock each operation and its output.

The function takes spark as a parameter. This needs to be mocked. Create created_table by calling spark.sql(create_table_query.format(**some_args)). We want to assert that spark.sql() is called only once. We’ll need to mock the output of spark.sql()as well. Coalesce created_table. Ensure that coalesce() is called with the parameter 1. Mock the output. Write the coalesced table. We need to mock .write, and mock the output of calling it on our table. Save the coalesced table to a save path. Ensure that it’s been called with the correct parameters.

As before, the test script is below, with comments for clarity.

Finally, save everything in a folder together. You could import the functions from their respective modules if you want to, or just have everything in one script.

To test it, navigate in the command line to your folder (cd folder/working_folder) and call:

python -m pytest final_test.py.

You should see something like (excuse the formatting):

serena@Comp-205:~/workspace$ python -m pytest testing_tutorial.py ============================= test session starts ============================== platform linux -- Python 3.6.4, pytest-3.3.2, py-1.5.2, pluggy-0.6.0 rootdir: /home/serena/workspace/Personal, inifile: plugins: mock-1.10.0 collected 3 items testing_tutorial.py ... [100%] =========================== 3 passed in 0.01 seconds ===========================

And that’s it!

There you have it. I hope this was somewhat helpful. When I was trying to figure out how to mock, I wish I had come across some sort of tutorial like this. Now go ahead, and as Stewie said, (don’t) stop mocking me (functions)!

Originally published at serena.mcdonnell.ca on November 7, 2018.