FrankWiles.com

Your pytest fixtures can return containers

As you probably know I’m a big fan of pytest fixtures and use them on nearly every project we work on at REVSYS.

I had a not uncommon situation that I solved in a new way today and wanted to share it. It’s VERY simple, but something I can imagine many developers over looking. I overlooked it until today.

The situation

I’m building a SaaS product with Django that has the usual problem of needing my User to have access to one or more Customer models. Like many SaaS products this one is going to have the pattern like Sentry, Rollbar, and many many others where as the currently logged in user you “switch” between which accounts you’re working with at any given moment.

In this product, Customer is our “account” model.

To do make this a little more sane I on the User model I’ve added a current_customer foreign key. When a user registers or gets invited to a customer this is set and it’s changed any time the user “switches”.

Pretty straight forward so far.

Normally, I would define a separate fixture for “Customer X” and also create a “User X” fixture. So I can do things like:

python
def test_switching(x_user, x_customer, y_customer):
  assert x_user.current_customer == x_customer

  switch(x_user, y_customer)

  assert x_user.current_customer == y_customer

This looks and seems logical, but there is a problem. These fixtures depend on each other. I can’t set x_user’s current customer without having the x_customer already built and while it’s probably possible to get the fixture ordering JUST right, it’s brittle.

Return a Container instead

Instead, we should return a small container object! I’m ommitting the imports for brevity. And yes, this could be a dataclass to make it even shorter.

python
class CustomerContainer:
    user = None
    customer = None

    def __init__(self, customer, user):
      self.customer = customer
      self.user = user

@pytest.fixture
def x(db):
  user, created = User.objects.get_or_create(...)
  customer, created = Customer.objects.get_or_create(...)

  user.current_customer = customer
  user.save()

  return CustomerContainer(customer=customer, user=user)

This allows us to build these together, pass around a small object to access them in a natural feeling way. Now our test can look like this.

python
def test_switching(x, y):
  assert x.user.current_customer == x.customer

  switch(x.user, y.customer)

  assert x.user.current_customer == y.customer

Obviously you shouldn’t be using x and y as your fixture names in favor of something more descriptive. Because I’m building this product with a friend named “Sam” I actually have the fixture containers named frank and sam to hold our two Customer objects and the associated users.

Headshot of Frank Wiles

Frank Wiles

Founder of REVSYS and former President of the Django Software Foundation . Expert in building, scaling and maintaining complex web applications. Want to reach out? Contact me here or use the social links below.