How do I test a method in Django which closes the database connection?

Question

I have a long running Python process which uses the Django ORM. It looks something like this:

import django
from django.db import connection

from my_app import models


def my_process():
    django.setup()
    while (True):
        do_stuff()
        connection.close()
        models.MyDjangoModel().save()

Sometimes do_stuff takes a really long time, at which point I was running into an error where my MySql connection timed out because the database server killed the connection as idle. Adding the connection.close() line forces django to get a new connection every time and fixes that issue. (See https://code.djangoproject.com/ticket/21597).

However, I am testing this process using a django.test.TestCase, and calling connection.close causes those tests to fail, as django's TestCase class wraps the test in a transaction, and closing the connection within that transaction causes the transaction to break and raises a django.db.transaction.TransactionManagementError.

One attempt at resolving this issue I tried is setting the CONN_MAX_AGE database parameter and calling connection.close_if_unusable_or_obsolete instead, but the transaction also changes the connection's autocommit setting from the settings' default value of True to False which in turn causes close_if_unusable_or_obsolete to try and close the connection anyway (https://github.com/django/django/blob/master/django/db/backends/base/base.py#L497).

I guess I could also mock connection.close in test so it does nothing, but that seems kind of hacky.

What is the best way to test a django method which needs to close the database connection?

Answer 1

Well, I'm not sure if this is the best answer, but since this question has seen no responses so far, I will post the solution I ended up using for posterity:

I created a helper function which checks whether we are currently in an atomic block before closing the connection:

import django
from django.db import connection

from my_app import models


def close_connection():
    """Closes the connection if we are not in an atomic block.

    The connection should never be closed if we are in an atomic block, as
    happens when running tests as part of a django TestCase. Otherwise, closing
    the connection is important to avoid a connection time out after long actions.     
    Django does not automatically refresh a connection which has been closed 
    due to idleness (this normally happens in the request start/finish part 
    of a webapp's lifecycle, which this process does not have), so we must
    do it ourselves if the connection goes idle due to stuff taking a really 
    long time.
    """
    if not connection.in_atomic_block:
        connection.close()


def my_process():
    django.setup()
    while (True):
        do_stuff()
        close_connection()
        models.MyDjangoModel().save()

As the comment states, close_connection prevents connection.close from being called under test, so we no longer break the test transaction.

Answer 2

Instead of modifying myprocess function you can inherit your Test class from TransactionTestCase which doesn't wrap your test function in transaction atomic block in the first place. If you are keen about having the transaction atomic in your test function then wrap it only around where you need it.

class Test(TransactionTestCase):
     def test_function(self):
         do_stuff()
         with transaction.atomic():
             your_transaction_atomic_test()
         my_process()
         do_other_stuff()

ps: Inheriting your Test class from TestCase will wrap your total test function with transaction atomic block.