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 datetime
import imported_classdef sample_test():
# Arrange
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
assert test_result
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 True or 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
assert not test_resultPython 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.
Fixtures
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:
@pytest.fixture
def setup_sample_test():
cursor.execute(
"""
INSERT INTO JOB_DB.test_job_history (
start_time,
end_time,
job_status,
input_start_date_hour,
input_end_date_hour
) VALUES (
'2020-05-22 10:00:00',
null,
'RUNNING',
'2020-01-31 21:00:00',
'2020-01-31 22:00:00'
)
"""
) yield
clean_up()
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 setup_sample_test: we 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
yieldstatement.
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?
Autouse Fixtures
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 pd@pytest.fixture(scope="session", autouse=True)
def db_conn():
snowtest.create()
conn = snowtest.create_db_connection() def execute(sql, parameters=None, **kwargs):
cur = conn.cursor()
cur.execute(sql, parameters, **kwargs)
return pd.DataFrame.from_records(
iter(cur), columns=[x[0] for x in cur.description]
).to_dict() global cursor
cursor = conn.cursor(DictCursor)
yield
snowtest.destroy()
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[0]). 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:
def sample_test(setup_sample_test):
# Arrange
...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.
Sources
- pytest fixtures: explicit, modular, scalable, pytest docs
- pandas.DataFrame.from_records, pandas docs
- Python iter(), Programiz
- What is the cursor execute in Python?, Shekhar Pandey, Linux Hint