Code Testing in Python: Red, Green and Refactor

Last week, we went over some tips and tricks with code typing (which can be found here). In this blog, we're going to go over some code testing processes with Python. Before we get started, here's a small bit of context. Let’s assume our system operates over logistics data grouped by date and zip code. It could look something like the following:

{
  "12345": [
    {"day": "2022-01-01", "deliveries": 1.123},
    {"day": "2021-12-31", "deliveries": 3.14},
    {"day": "2021-12-30", "deliveries": 1.618},
    {"day": "2021-12-29", "deliveries": 6.022},
    ...,
  ],
  "23456": [...],
  ...
}

All has been working well for logistics on the West Coast, but processing data for the rest of the US has caused production issues. We have tracked a series of 500 errors to their source, but our logs don’t provide much detail. Here’s the relevant code snippet:

# daily/report.py
from datetime import date

def deliveries_by_day(logistics, zips, day=date.today()):
  deliveries = 0
  for zip_code in zips:
	  for data in logistics[zip_code]:
	    if data["day"] == day:
		  deliveries += data["deliveries"]
		  break

  return deliveries

...
# logistics = {"12345": [{...}, ...], ...}
# CT_HARTFORD = {"06101", "06102", ...}
report = deliveries_by_day(logistics, CT_HARTFORD)
...

On the surface, all looks as expected: zip codes for Hartford, CT, are defined as a set of strings from which today’s count of deliveries are summed up. However, command line (i.e., manual) testing shows that python is throwing `KeyError: ‘06101’` lookup exceptions. Examining the logistics data, we find that zip code keys have all been trimmed when there is one or more leading zero, i.e.:

print(logistics)
"""
{
	...,
	"6101": [...],  # 06101
	"6102": [...],  # 06102
	"501": [...],   # 00501
	...
}
"""

Our code is fixed, validated with a test, and ready to commit! But we have accidentally discovered a potential issue while writing our test: the code fails if there isn’t a “deliveries” key in the logistics dictionary. After a short discussion, the team decides to resolve this issue at the same time. Research into the logistics data shows that the “day” key is always present, but all others appear to be optional. It would be simple to just fix the code as part of our refactor, but this specific edge case deserves a dedicated (aka “unit”) test.

# tests/test_deliveries_by_day.py
from daily.report import deliveries_by_day

logistics = {
  ...
  "12345": [{"day": "2022-01-01", "foo": "bar"}],
}

...

def test_deliveries_key_not_present():
  assert deliveries_by_day(logistics, {"12345"}, "2022-01-01") == 0

The test fails as expected, so it’s time to make minimal changes so that the test passes:

def deliveries_by_day(logistics, zips, day=date.today()):
   return sum(
     [
-      data["deliveries"]
+      data.get("deliveries", 0)
       for zip_code, days in logistics.items()
       if zip_code in [str(int(x)) for x in zips]
       for data in days

Now, we have both of our dedicated (unit) tests passing:

$ pytest -v
...
tests/test_deliveries_by_day.py::test_zip_starts_with_0 PASSED [ 50%]
tests/test_deliveries_by_day.py::test_deliveries_key_not_present PASSED [100%]
2 passed in 0.01s