How Reliable are your Queries?

Unit Testing for SQL

When working with data, it’s easy to forget that the queries we’re writing are code. Why should they be held to any less of a standard than what we write in our other projects? My team has achieved a test-driven development process for our queries with dbt and Snowflake — I’m going to show you how.

The Traditional Test-Driven Loop

The usual test-driven process has three main steps:

1. Write a failing test
2. Add code to make the test pass
3. Refactor, without making the test fail again

The benefit of writing failing tests as the first step is twofold: the process of writing a test requires us to fully understand the requirements of our code.
That is, the test can serve as a pseudo-specification of the external behavior of what we’re about to write. While tests don’t dictate how we implement
behaviors, they can dictate what the behaviors need to be.
The second benefit is that we can now prove that any code we write is correct. Code that passes our tests will have been shown to meet our specification
(or at least, whatever the test specifies).
So what does this have to do with SQL? Well, the queries we write are code that describes some data transformation — we might as well treat it like any other code.

The SQL Unit Test

The structure of SQL queries does not lend itself to unit testing: the transformation is tightly coupled to the input data (due to our FROM clause) and there’s not much we can do to inject fake data (this practice of mocking components that aren’t being tested is common in unit tests).

An Idea of a Test

There are certainly ways to verify the behavior of a SQL query by looking at the resulting rows. In the broader sense of a “test”, we really care about two
things: what goes in, and what comes out. By comparing what we get out of our query to a known “expected” output, we can ensure our query works as
we want it to. Knowing this, we just need a way to express this idea as SQL.

Enter dbt

dbt is an open source tool for data transformations that works alongside most data warehouse solutions. In short, it is the “Transform” of the “Extract,
Load, Transform” pipeline.
On top of being a fantastic tool for building out series of data transformations, dbt allows the use of Jinja, a templating language built on top of Python.
Jinja lets us insert some standard programming structures into SQL, so long as the final compiled query is only SQL. For example, we can use a for loop
and a template to quickly build SQL queries:

{% set payment_methods = ["bank_transfer", "credit_card", "gift_card"] %}
 select
order_id,
{% for payment_method in payment_methods %}
sum(case when payment_method = '{{payment_method}}' then amount end) as {{payment_method}}_amount,
{% endfor %}
sum(amount) as total_amount
from app_data.payments
group by 1

This templating gets compiled into SQL like this:

select 
  order_id, 
  sum(case when payment_method = 'bank_transfer' then amount end) as bank_transfer_amount, 
  sum(case when payment_method = 'credit_card' then amount end) as
credit_card_amount, 
  sum(case when payment_method = 'gift_card' then amount end) as gift_card_amount, 
  sum(amount) as total_amount 
from app_data.payments 
group by 1

While this is useful in building our queries in general, we care more about using this for testing. There are a couple Jinja macros that we are particularly interested for this purpose.

select
...
from {{ ref('my_other_model') }}
select
...
from {{ source('my_schema', 'my_table') }}

These two macros (ref and source) are used to build references to other dbt models or dbt data sources. While they serve as the main way to create
links between models, they’re also a potential injection point for mocked data (solving one of the biggest problems with SQL unit testing, the coupling of a
query with its input data).

By overriding these macros, we can replace the real data with our fake data when in the context of a unit test. The implementation of this is a bit out of
scope for this article, but the details can be seen on the Github repo for the unit testing package.

We’ve Got the Pieces, Now What?

At this point we know what a test should “look like” (provide data in, expect data out) and dbt provides us an injection point where we can inject our fake
data. Let’s write a test!

As an example, we’ll use the “Jaffle Shop” demo. The Jaffle Shop is a running example in dbt documentation, and these test cases are included in the unit
testing package source.

Suppose we have a model, customer_stats, that stores some information about our customers. We also have two models on which this model
depends: customers and orders. Our test is testing to make sure our customer statistics has the proper total lifetime value for each customer. Here’s what our test looks like:

{% call dbt_unit_testing.test('customer_stats', 'should sum order values to calculate customer_lifetime_value') %}
  
  {% call dbt_unit_testing.mock_ref ('customers') %}
    select 1 as customer_id, 'John' as first_name, 'Doe' as last_name
  {% endcall %}
  
  {% call dbt_unit_testing.mock_ref ('orders') %}
    select 1001 as order_id, 1 as customer_id, 10 as price
    UNION ALL
    select 1002 as order_id, 1 as customer_id, 10 as price
  {% endcall %}
{% call dbt_unit_testing.expect() %}
    select 1 as customer_id, 20 as customer_lifetime_value
  {% endcall %}
{% endcall %}

All of these are Jinja macros: after compilation we’re left with a query that runs the model we’re testing with our mock data and compares the result with the expected output. dbt has built-in data tests which fail when a query returns any rows, so the expect macro uses a couple EXCEPT SQL clauses to find all the mismatched rows between our output and the target.

Let’s break this down.

{% call dbt_unit_testing.test('customer_stats', 'should sum order values to calculate customer_lifetime_value') %}

Any line wrapped in {% … %} is a Jinja macro, meaning it executes some code in our templating engine. Here we’re calling the dbt_unit_testing.test macro, which takes in the name of the model we’re testing (customer_stats) and a description ("should sum order values to calculate customer_lifetime_value”).

{% call dbt_unit_testing.mock_ref ('customers') %}
  select 1 as customer_id, 'John' as first_name, 'Doe' as last_name
{% endcall %}

Here we’re defining a mock of a referenced model called customers. This is one of the models our customer_stats model depends on. Whenever we encounter a reference to this model, dbt will replace it with a CTE containing the SQL we wrap in the call block. This is our main method of injecting mocked data into our tests. When customer_stats tries to get a reference onto the real customers table in our database, it will instead be served this single row.

{% call dbt_unit_testing.mock_ref ('orders') %}
  select 1001 as order_id, 1 as customer_id, 10 as price
  UNION ALL
  select 1002 as order_id, 1 as customer_id, 10 as price
{% endcall %}

This is another mocked reference, this time for the orders model. We’ve defined two rows of fake data here, since we’re testing the sum of our prices.

{% call dbt_unit_testing.expect() %}
  select 1 as customer_id, 20 as customer_lifetime_value
{% endcall %}

Our expect call defines what we expect our outcome to be when running our model with the provided fake data. Here we’re expecting to see a customer_lifetime_value of 20, because they have two orders worth 10 each (defined in our mock above).

And that’s it! Defining unit tests with dbt is super quick and allows for proper test-driven development for SQL queries. Since we’ve implemented this as part of our workflow we have much more confidence in our queries and are able to validate their behavior before running them against massive datasets.