Database isolation in tests part 3

Database isolation in Django’s tests

In this post we’ll take a look at how Django handles test isolation. Wait, handles what? Isolation basically means that changes from one test don’t leak into another, so you know, just a fancy word for something we’ve been figuring out in last 2 parts. Same as in the last part, all cited source code will be modified to contain only important parts.

Important piece of knowledge is what database transaction is. I’ll do a very quick recap, but for more info you can go back to post about SQLAlchemy sessions. Also in this post we will assume database we use supports transactions and savepoints. DBs like postgres, mariadb, mssql all support both.

Part 1: click
Part 2: click

Transactions and rollback

Transactions are SQL specific, so it is a mechanism provided by database, not by Django, SQLAlchemy or any other python package. Transaction is like a bag for queries. At the end of transaction we can save all changes or discard them. The most important thing is all actions either succeed, fail or are discarded altogether. Only after transaction is committed (saved) data is saved to the database.

Let me bring back an example from a previous post.

BEGIN; -- transaction starts with BEGIN
-- all the queries you run are part of the transaction
SELECT * FROM users WHERE id = 1;

INSERT INTO orders (customer_id, order_date)
VALUES (1, NOW()) RETURNING id;

INSERT INTO order_items (order_id, product_id, quantity)
VALUES (1, 101, 2);

-- imagine calling an external service here
-- ROLLBACK transaction if something goes wrong.

UPDATE inventory SET stock = stock - 2 WHERE product_id = 101;

COMMIT; -- transaction ends with COMMIT or ROLLBACK

Savepoints

Transaction savepoints are marks in current transaction that we can rollback to. For example in case of error or if other logic forces us to undo only some of the changes.

BEGIN;
    INSERT INTO table1 VALUES (1);
    SAVEPOINT my_savepoint;
    INSERT INTO table1 VALUES (2);
    ROLLBACK TO SAVEPOINT my_savepoint;
    INSERT INTO table1 VALUES (3);
COMMIT;

BEGIN;
    INSERT INTO table1 VALUES (3);
    SAVEPOINT my_savepoint;
    INSERT INTO table1 VALUES (4);
    RELEASE SAVEPOINT my_savepoint; -- savepoint can be destroyed
COMMIT;

One thing important to notice, savepoints don’t give us ability to commit transaction multiple times. When we call COMMIT whole session gets committed, not the last savepoint.

Django’s TestCase

It is always the hardest to take a first step, right? Let me help you then. Why not start with the Django testing docs? Look at the first piece of code there.

from django.test import TestCase
from myapp.models import Animal


class AnimalTestCase(TestCase):
    def setUp(self):
        Animal.objects.create(name="lion", sound="roar")
        Animal.objects.create(name="cat", sound="meow")

    def test_animals_can_speak(self):
        """Animals that can speak are correctly identified"""
        lion = Animal.objects.get(name="lion")
        cat = Animal.objects.get(name="cat")
        self.assertEqual(lion.speak(), 'The lion says "roar"')
        self.assertEqual(cat.speak(), 'The cat says "meow"')

How do we know this is a good place to start? This test creates some database entries in setup and has access to them in test. We know that Django cleans up database between every test, and there’s no other setup required mentioned by docs. All “magic” has to be happening inside TestClass.

Before we delve into the magic world of Django internals we have to address a few things first so we don’t feel lost later.

TestCase, setUp and tearDown

When looking at Django test case we can clearly see that it is based on Python’s unittest package. Quick recap (or take a look in docs):

TestCase - used for aggregating multiple tests to one test case
setUpClass() - classmethod that runs once before all tests in the test case
TearDownClass() - classmethod that runs once after all tests in the test case
setUp() - method that runs before every test
TearDown() - method that runs after every test

Looking at example above, setUp method is available for a developer and you don’t have to call super() to trigger some extra logic. It looks like you have full control of what happens before and after test. It seems impossible for this example to work because there is no place it can clean db. But it does.

class AnimalTestCase(TestCase):
    def setUp(self):
        Animal.objects.create(name="lion", sound="roar")

    def test_animal_count1(self):
        Animal.objects.create(name="cat", sound="meow")
        self.assertEqual(Animal.objects.count(), 2)

    def test_animal_count2(self):
        Animal.objects.create(name="cat", sound="meow")
        self.assertEqual(Animal.objects.count(), 2)

Found 2 test(s).
..
----------------------------------------------------------------------
Ran 2 tests in 0.004s

OK

There has to be some other mechanism that allows for extra logic before setUp() and after tearDown(). If we go up inheritance tree to SimpleTestCase we’ll find it overwrites __call__ function.

 unittest.TestCase
    └── SimpleTestCase
        └── TransactionTestCase
            └── TestCase

def __call__(self, result=None):
    """
    Wrapper around default __call__ method to perform common Django test
    set up. This means that user-defined TestCases aren't required to
    include a call to super().setUp().
    """
    self._setup_and_call(result)

def _setup_and_call(self, result, debug=False):
    """
    Perform the following in order: pre-setup, run test, post-teardown
    """
   self._pre_setup()

    if debug:
        super().debug()
    else:
        super().__call__(result)

    self._post_teardown()

So Django overwrites __call__ to inject _pre_setup and _post_teardown methods for its own use so you can use setUp and tearDown yourself without calling super(). That’s very neat!

But why __call__? What does it give us? Also it is not really clear how this logic is run for every test, not just once the test case. In depth explanation will take us away too far from the main topic. Let’s keep it very simple, even oversimplified. You can run test by calling test case instance.

TestCase.__call__(result=None) runs one test, and returns result that is passed to next test.

result = None
for test_name in TestCase.get_all_test_names():
    test_case = TestCase(test_name)
    result = test_case(result)

return result

Again, this is huge oversimplification and not how things work under the hood, but important thing to remember is test case instance is called once per test, not per test case.

After you’re finished with this post I encourage you to look at unittest docs for more info.

How django cleans database up after each test?

This is what we’ve all been waiting for. Before we dive deep into logic quick reminder:

• database transaction is a bag for SQL queries which results are visible inside current transaction and saved or discarded at the end
• savepoint allows us rollback (discard) only a part of transaction
• thanks to architecture of Django and its test client tests and application can and do share same database connection
• in Django, fixtures are extra data loaded to database before tests. Used for populating db with things like constant values or user account used in all tests.

Class setup and teardown.

Let’s start with TestCase to check if it does any setup for whole test case.

class TestCase(TransactionTestCase):

    @classmethod
    def setUpClass(cls):
        super().setUpClass()

        cls.cls_atomics = cls._enter_atomics()

        if cls.fixtures:
            for db_name in cls._databases_names(include_mirrors=False):
                try:
                    call_command("loaddata", *cls.fixtures, verbosity=0, database=db_name)
                except Exception:
                    cls._rollback_atomics(cls.cls_atomics)
                    raise

        # omitted: test data setup, and TestData wrapper

class TestCase(TransactionTestCase):

    @classmethod
    def _enter_atomics(cls):
        """Open atomic blocks for multiple databases."""
        atomics = {}
        for db_name in cls._databases_names():
            atomic = transaction.atomic(using=db_name)
            atomic._from_testcase = True
            atomic.__enter__()
            atomics[db_name] = atomic
        return atomics

Django uses @transaction.atomic (docs) decorator to wrap piece of code in transaction. setUpClass() creates instance the same atomic class (although not used as decorator). So each test runs inside a transaction! Keep in mind that it loads fixtures after calling _enter_atomics() .

So if setUpClass() starts transaction, teardown should close it. Let’s check that quickly.

class TestCase(TransactionTestCase):

    @classmethod
    def tearDownClass(cls):
        cls._rollback_atomics(cls.cls_atomics)
            for conn in connections.all(initialized_only=True):
                conn.close()
        super().tearDownClass()

class TestCase(TransactionTestCase):

    @classmethod
    def _rollback_atomics(cls, atomics):
        """Rollback atomic blocks opened by the previous method."""
        for db_name in reversed(cls._databases_names()):
            transaction.set_rollback(True, using=db_name)
            atomics[db_name].__exit__(None, None, None)

Yes! It marks transaction to rollback and exits atomic that was set up in setUpClass. Let’s sum up what is going on in TestCase so far:

1. setUpClass starts transaction for whole test case. Keeps related Atomic in cls.atomics.
2. loads fixtures
3. runs tests
4. marks transaction to be rolled back and exits Atomic from point 1

_pre_setup and _post_teardown

We also know there is extra per-test logic inside _pre_setup and _post_teardown. Let’s check it.

class TransactionTestCase(SimpleTestCase):  # is parent of TestCase

    @classmethod
    def _pre_setup(cls):
        """
        Perform pre-test setup:
        * If the class has a 'fixtures' attribute, install those fixtures.
        """
        super()._pre_setup()
        # omitted: set installed apps
        try:
            cls._fixture_setup()
        except Exception:
            # omitted: unset installed apps
            raise

        # omitted: clear query log

class TestCase(TransactionTestCase):

    @classmethod
    def _fixture_setup(cls):
        # omitted: handling lack of transaction support

        cls.atomics = cls._enter_atomics()

        # omitted: handling lack of savepoints support

class TestCase(TransactionTestCase):

    def _fixture_teardown(self):
        # omitted: handling lack of transaction support
        # omitted: constraints check

        self._rollback_atomics(self.atomics)

class TransactionTestCase(SimpleTestCase):  # is parent of TestCase

    def _post_teardown(self):
        self._fixture_teardown()
        super()._post_teardown()

        # omitted: close db connections
        # omitted: handle installed apps

As we can see per-test logic is really similar to logic for whole test case. Only difference is we don’t load fixtures every test, which makes sense. We have to look at Atomic implementation.

Django’s Atomic

I have to address few things now. Firstly, Django's Atomic sounds like crossover between Tarantino’s movie and Fallout - 100% would watch. Second is we are digging scarily deep inside Django internals. I remember myself using @transaction.atomic few years back and thinking “oh my, I don’t even want to think how complicated it is under the hood. Hope it doesn’t break some day, because it would take me weeks to wrap my head around that”. Turns out, I couldn’t have been further from the truth.

Most developers probably never had to think about that, but Django’s default behavior is autocommit (source). What that means? Without using @transaction.atomic you get experience similar to using your database CLI. Each query resulting in changes is immediately saved to DB without ability to undo. Remember last time you updated too many rows because you forgot WHERE clause? Yeah, that behavior.

On the other hand with autocommit=False database driver implicitly starts new transaction by injecting BEGIN before your query. More on that can be found at the top of docs for postgres driver and in autocommit section

If you’re not interested in looking at source code anymore, docs has explanation good enough to understand this double __enter__ and __exit__. Here: it is pretty short.

Deep dive

All logic resides here. Let’s start from the top, there are a bunch of functions, and are pretty self explanatory:

get_connection() - get underlying database connection object
get|set_autocommit() - get/set autocommit behavior we talked about minute ago
commit() - calls commit
rollback() - calls rollback
savepoint() - creates savepoint in transaction
savepoint_rollback() - rollback to savepoint
savepoint_commit() - releases savepoint
set_rollback() - marks db connection (transaction) for rollback

class Atomic(ContextDecorator):

    def __enter__(self):
        connection = get_connection(self.using)
        # omitted: durable atomics

        if not connection.in_atomic_block:
            # Reset state when entering an outermost atomic block.
            connection.commit_on_exit = True
            connection.needs_rollback = False
            # omitted: connection's autocommit = false

        # omitted: need rollback logic
        if connection.in_atomic_block:
            sid = connection.savepoint()
            connection.savepoint_ids.append(sid)
        else:
            connection.set_autocommit(
                False, force_begin_transaction_with_broken_autocommit=True
            )
            connection.in_atomic_block = True

        # omitted: durable atomics

If we are inside first Atomic context manager:

1. mark connection to commit transaction on exit and don’t roll it back
2. set connection autocommit to False and mark we are in atomic block - that will implicitly start transaction with next query

Next time we use context manager or __enter__ our connection is already marked with in_atomic_block, so:

1. create a savepoint and append it to connection’s savepoints list

We have pretty good grasp what happens when we enter Atomic. Let’s look at exit now.

class Atomic(ContextDecorator):
    def __exit__(self, exc_type, exc_value, traceback):
        connection = get_connection(self.using)

        # omitted: durable atomics
        # omitted: all error handling
        # omitted: savepoint releasing
        # changed flow for easier read

        # if we don't have any savepoint, we are exiting whole transaction
        if connection.savepoint_ids:
            sid = connection.savepoint_ids.pop()
        else:
            connection.in_atomic_block = False

        if not connection.in_atomic_block:
            if connection.needs_rollback:
                connection.rollback()
            else:
                connection.commit()
        else: # still in atomic block
            if connection.needs_rollback:
                connection.savepoint_rollback(sid)

        if not connection.in_atomic_block:
            connection.set_autocommit(True)

I heavily redacted source code so it is easier to follow. Let’s sum it up

1. If we don’t have any savepoints, we are outermost Atomic, so after exit we are no longer in atomic block.
2. If we are no longer in atomic block, we either commit or rollback whole transaction.
3. Else, if connection is marked to rollback we rollback only to the last savepoint.
4. If no longer in atomic block, set autocommit back to True so further queries are not wrapped in transaction anymore.

So pretty much arrived at what Django documentation says. That was pretty intensive, I think we need summary of Atomic inside Atomic:

with transaction.atomic():  # outer atomic
    # set connection's autocommit to False
    # implicitly start a new transaction
    ...

    with transaction.atomic(): # inner atomic block
        # create a savepoint
        ...

    # exit inner atomic block, rollback or commit savepoint

# exit outer atomic, rollback or commit transaction
# set connection's autocommit to True

Tying it all together

We have all pieces in place to fully understand how Django cleans database between tests.

1. When TestCase sets up test case it creates new transaction by entering transaction.atomic and loads fixtures.
2. Before each test TestCase enters transaction.atomic again, to create savepoint in transaction.
3. After each test TestCase marks transaction to rollback and exits transaction.atomic. It rolls back transaction to the savepoint before test and thus cleans all changes made during test.
4. After all tests are done TestCase exits transaction.atomic and rolls back transaction it started during test case setup. It cleans loaded fixtures.

This is all possible because database connection is shared across tests, view functions and any other place. Also transaction.atomic can be nested without committing or rolling back whole session - just to a savepoint. As you can see there are multiple architecture decisions, that allowed to create abstractions to seamlessly deliver database cleanup by wrapping it in transaction. Transaction is almost never explicitly committed in Django - there is very little risk of committing transaction that was meant to be rolled back at the end of tests. Fact that ORM is integral part of whole framework helps immensely to deliver such solution that work out of the box.

It is ideal, best solution? Not really. This approach won’t let you test any behaviors related to committing transaction. For this purpose Django provides TransactionTestCase which uses different mechanism for cleanup. You can also argue that by using TestCase you test different behavior than in production setup. I would personally agree, but this approach works for thousands of developers so it seems to be a pretty sane default.

Like many other things in programming this is a tradeoff. This time between tests speed and test setup that doesn’t ideally reflect production conditions.

Summary

If you are still here - thank you! That wasn’t easy.

You probably noticed that I bluntly stated that database connection is shared everywhere in Django and never showed that explicitly? Well, this is another journey. But we brushed against it when we looked at atomic source code. It gets connection objects from global django.db.connections. But for now I’ll leave that. You can take a shot and dig into source yourself.

Now we have pretty good insight on how Django manages database transactions in tests. We also now understand that many pieces have to fall in right places to enable transaction wrapping in tests. Even with all that, this solution isn’t perfect.

In next part we will look at what options we have with FastAPI and SQLAlchemy and do some benchmarks. We’ll also show why Django’s approach is not easy to recreate there. Don’t worry, we are not leaving TransactionTestCase for good. We will go through its cleanup mechanisms in relation to other frameworks.

See ya!