Loyal readers might be familiar with a certain Python script in my workplace. If not, here’s what you need to know: it’s super long and hard to edit because it has so many things going on. Thanks to some additional resources and oversight now being provided to my team, we’re doing a lot of work on this script to solve those problems. Part of that work is actually writing tests for this thing.
Off the bat, I was given the recommendation to investigate whether we had a company-wide style guide for testing in Python. While I didn’t find that, I did get some very helpful advice from a Senior Engineer, who showed me some examples of tests his team has developed. I was very grateful and went about reading the files. After a few minutes, I realized that I’d have to do some background work on
pytest because I was having a lot of trouble understanding.
A test in pytest
import imported_classdef sample_test():
current_datetime = datetime.datetime.strptime(
"2020-05-22 13:00:00", "%Y-%m-%d %H:%M:%S"
) # Act
test_result = imported_class.function_to_test(current_datetime) # Assert
As usual, here’s a simplified version of production code. The actual test has a complicated name, but it doesn’t do much more than what we’re seeing here. Thanks to knowledge from a previous post, I knew that Arrange was my setup step and Act would execute the functionality we want to test. In this case, we’re creating a timestamp and then passing it to our test function. In the production code, it actually executes a function that returns
False depending on whether a job will run at a certain time of day. We Assert that the result will be
True because we’re testing that the job will run at the time of day stated. In another test, we can use a different timestamp and change our assertion:
assert not test_result
Python lives up to its reputation as an “easy to read” language here. We’ve written a test without having to learn any new syntax. The complexity of how the test logic works is abstracted within our
function_to_test function, so there’s nothing too difficult to understand here. It’s definitely possible to put that language directly into the test, but it might rely on other functions or variables so it probably makes the most sense to stay in the actual script module.
If reading a test file was this easy, I wouldn’t have had to write a blog post about it. I decided on this topic because before I even got to see a test like this, I saw a different kind of function:
INSERT INTO JOB_DB.test_job_history (
) VALUES (
According to the pytest docs, test fixtures initialize test functions…Initialization may setup services, state, or other operating environments. In other words, fixtures are setup functions that prepare the environment for a test to run. They are denoted with the
@pytest.fixture flag and are passed to tests as arguments.
Here we see that the function’s main purpose is to add a row to our
test_job_history table for a job whose status is
RUNNING. If we use this in conjunction with our previous test, we can check whether a new job will run while an existing job is already running. In production, this is specifically meant to check that a new job will not start if another job is still running and it’s been fewer than 3 hours since that first one kicked off. It’s acomplicated condition that we want to make sure to get right.
Two other things happen in
yield and we
clean_up. The same pytest docs page explains that yield fixtures are preferred to using
return to get a result back:
Once pytest figures out a linear order for the fixtures, it will run each one up until it returns or yields, and then move on to the next fixture in the list to do the same thing.
Once the test is finished, pytest will go back down the list of fixtures, but in the reverse order, taking each one that yielded, and running the code inside it that was after the
What comes after the
yield statement? Our
clean_up function! By following this pattern, we can return data generated by our fixtures but then safely remove it after the test has finished so we clog up our memory. In this case, we’re still following the pattern even though our function doesn’t actually return anything. Its only responsibility is to update an existing DB. Speaking of which, how can it use
cursor if that hasn’t been defined to this point?
Near the top of the test file I was reading, I found a fixture that was using a couple of arguments:
from test_utils.snowtest.snowtest import Snowtest
from snowflake.connector import DictCursor
import pandas as email@example.com(scope="session", autouse=True)
conn = snowtest.create_db_connection() def execute(sql, parameters=None, **kwargs):
cur = conn.cursor()
cur.execute(sql, parameters, **kwargs)
iter(cur), columns=[x for x in cur.description]
).to_dict() global cursor
cursor = conn.cursor(DictCursor)
This is a little complicated, but our main takeaway is that the function creates a
global cursor which we later use to make DB edits. Its annotation sets the scope as
session and sets
autouse as True. That means that this is going to be run even if we don’t explicitly reference it in our test.
To break down the logic, we’re importing the
snowflake.connector library that allow us to query Snowflake databases and a custom
snowtest library from our own utils. The
snowtest import allows us to create a blank Snowflake DB with a custom series of tables (this logic is done outside of the function, but we just pass our table names as strings).
conn is our connection to that test DB.
We define a function,
execute, which we have already seen being invoked in
setup_sample_test. It runs a SQL command and then uses pandas.DataFrame.from_records to turn the resulting data into a DataFrame. The
from_records function accepts an iterable as its first argument, in this case we’re transforming our cursor object into an iterable with iter. The cursor object includes all the data we’ve just added to our test DB. It can also accept an optional
columns argument, which determines how to name the columns in the resulting DataFrame. Here, we iterate through the elements in
cur.description and take out the name from each (
x). Finally, we turn that DataFrame into a dictionary and return that value. This is a cool customization because the execute function, by default, just returns an integer of how many rows were affected by our query. This custom execute gives us a lot more data in the return.
Next we create our global cursor, and connect it with our test Snowflake DB. After that, we yield and clean up in a pattern we’re now familiar with.
Bringing it all together
Because we added the
autouse=True argument to
db_conn, we don’t have to explicitly invoke it before any of our tests. That’s not the case for
setup_sample_test, however, so we would pass it in to our very first function,
sample_test as an argument:
This gave me a really critical understanding of how to integrate Snowflake tables and queries into my pytest functions. The script I’m hoping to test interacts heavily with DBs and I wasn’t sure how I would make the connection or if the tests would be expensive because of running heavy queries. It’s a relatively simple example, but in my experience, getting just a simple version of your code to run is among the biggest hurdles in programming.