The problem

At this point we had a large codebase and a lot of tests. Unfortunately, a lot of our tests were a relic from when we were using fixtures extensively.

The thought of having to update and add new fixtures was unacceptable. But, we still had to test the new features so we started writing tests like this:

def test_should_charge_credit_card(self):

feature_set = user_account.feature_set

feature_set.can_pay_with_credit_card = True

feature_set.save(update_fields=['can_pay_with_credit_card'])

pay_with_credit_card(user_account, 100)



def test_should_fail_when_feature_disabled(self):

feature_set = user_account.feature_set

feature_set.can_pay_with_credit_card = False

with self.assertRaises(FeatureDisabled):

pay_with_credit_card(self.user_account, 100)

We had a lot of tests to update and some of the features we added interrupted the flow of other tests which resulted in a mess!

The context manager

We already used context managers to improve our tests in the past, and we thought we can use one here to set features on and off:

from contextlib import contextmanager @contextmanager

def feature(feature_set, feature_name, enabled):

original_value = getattr(feature_set, feature_name)

setattr(feature_set, feature_name, enabled)

feature_set.save(update_fields=[feature_name]) def feature(feature_set, feature_name, enabled):original_value = getattr(feature_set, feature_name)setattr(feature_set, feature_name, enabled)feature_set.save(update_fields=[feature_name]) try:

yield finally:

setattr(feature_set, feature_name, original_value)

feature_set.save(update_fields=[feature_name])

What does this context manager do?

Save the original value of the feature. Set the new value for the feature. Yields — this where our test code actually executes. Set the feature back to the original value

This made our tests much more elegant:

def test_should_charge_credit_card(self):

with feature(

user_account.feature_set,

can_pay_with_credit_card,

True,

):

pay_with_credit_card(user_account, 100) def test_should_fail_when_feature_disabled(self):

with feature(

user_account.feature_set,

can_pay_with_credit_card,

False,

):

with self.assertRaises(FeatureDisabled):

pay_with_credit_card(self.user_account, 100)

This is was a big step forward but we were still not satisfied.

**kwargs

This context manager has proven to be very useful for features so we thought… why not use it for other things as well?

We had a lot of methods involving more than one feature:

def test_should_not_send_notification(self):

feature_set = user_account.feature_set

with feature(feature_set, can_pay_with_credit_card, True):

with feature(feature_set, can_receive_notifications, False):

pay_with_credit_card(user_account, 100)

Or more than one object:

def test_should_not_send_notification_to_inactive_user(self):

feature_set = user_account.feature_set

user_account.user.is_active = False

with feature(feature_set, can_receive_notifications, False):

pay_with_credit_card(user_account, 100)

So we rewrote the context manager to accept any object and added support for multiple arguments:

@contextmanager

def temporarily(obj, **kwargs):

original_values = {k: getattr(obj, k) for k in kwargs} def temporarily(obj, **kwargs):= {k: getattr(obj, k) for k in kwargs} for k, v in kwargs.items():

setattr(obj, k, v) obj.save(update_fields=kwargs.keys()) try:

yield finally:

for k, v in original_values.items():

setattr(obj, k, v) obj.save(update_fields=original_values.keys())

The context manager can now accept multiple features, save the original values, set the new values and restore when we are done.

Testing became much easier:

def test_should_not_send_notification(self):

with temporarily(

user_account.feature_set,

can_pay_with_credit_card=True,

can_receive_notifications=False,

):

pay_with_credit_card(user_account, 100)

self.assertEquals(len(outbox), 0)

We can now use the function on other objects as well:

def test_should_fail_to_login_inactive_user(self):

with temporarily(user, is_active=False):

response = self.login(user)

self.assertEqual(response.status_code, 400)

Profit!

The hidden performance benefit

After a while getting comfortable with the new utility we noticed another performance benefit.

In tests that had heavy setups we managed to move the setup from the test level to the class level.

To illustrate the difference let’s test a function that sends an invoice to the users. Invoices are usually sent only when the transaction is complete. To create a complete transaction we need a lot of setup (choose products, checkout, issue payment etc).

This is a test that require a lot of setup:

class TestSendInvoice(TestCase): def setUp(self):

self.user = User.objects.create_user(...)

self.transaction = Transaction.create(self.user, ...)

Transaction.add_product(...)

Transaction.add_product(...)

Transaction.checkout(...)

Transaction.request_payment(...)

Transaction.process_payment(...)



def test_should_not_send_invoice_to_commercial_user(self):

self.user.type = 'commercial'

mail.outbox = []

Transaction.send_invoice(self.user)

self.assertEqual(len(mail.outbox), 0)



def test_should_attach_special_offer_to_pro_user(self):

self.user.type = 'pro'

mail.outbox = []

Transaction.send_invoice(self.user)

self.assertEqual(len(mail.outbox), 1)

self.assertEqual(

mail.outbox[0].subject,

'Invoice and a special offer!'

)

The setUp function need to execute before each test function because the test functions change the objects and that might create a dangerous dependency between test cases.

To prevent dependencies between test cases we need to make sure each test leaves the data exactly as it got it. Luckily, this is exactly what our new context manager does:

class TestSendInvoice(TestCase): @classmethod

def setUpTestData(cls):

cls.user = User.objects.create_user(...)

cls.transaction = Transaction.create(cls.user, ...)

Transaction.add_product(...)

Transaction.add_product(...)

Transaction.checkout(...)

Transaction.request_payment(...)

Transaction.process_payment(...)



def test_should_not_send_invoice_to_commercial_user(self):

mail.outbox = []

with temporarily(self.user, type='commercial'):

Transaction.send_invoice(self.user)

self.assertEqual(len(mail.outbox), 0)



def test_should_attach_special_offer_to_pro_user(self):

mail.outbox = []

with temporarily(self.user, type='pro'):

Transaction.send_invoice(self.user)

self.assertEqual(len(mail.outbox), 1)

self.assertEqual(

mail.outbox[0].subject,

'Invoice and a special offer!'

)

We moved the setup code to setUpTestData. The setup code will execute only once for the entire test class resulting in quicker tests.

Final words

The motivation for this context processor was our long unhealthy relationship with fixtures. As we scaled our app the fixtures became a burden. Having so many tests rely on them made it difficult to completely replace.

With the addition of features we knew we did not want to rely on fixtures any more and we looked for creative, more verbose and maintainable ways, of managing test data. Having a simple way to create different variations of an object for testing was exactly what we needed.