It all started when we added feature flags to our app. After some deliberation we created a "feature set" model with boolean fields for each feature:
class FeatureSet(models.Model):
name = models.CharField(max_length=50)
can_pay_with_credit_card = models.BooleanField()
can_save_credit_card = models.BooleanField()
can_receive_email_notifications = models.BooleanField()
We added a foreign key from the user account to the feature sets model, and created feature sets for "pro", "newbie" and "commercial" users.
To enforce the features we added tests in appropriate places. For example:
def pay_with_credit_card(self, user_account, amount):
if not user_account.feature_set.can_pay_with_credit_card:
raise FeatureDisabled('can_pay_with_credit_card')
# ...
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])
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)
**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}
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.
