reformy
reformy

Reputation: 1027

Django pytest: Clear DB after a running a test case

I have a test (django pytest) that needs to manipulate objects in DB. The thing is that after the test, the DB is "dirty" and other tests fail. I saw something about TransactionTestCase but I can't understand how it works with django test.

Here is a simple example of my current code:

@pytest.mark.django_db
def test_something(mock_my_obj):
    mock_my_obj.save()
    # test test test
    ...
    # I don't want to delete the obj here...

UPDATE: Second try: I read that TestCase should use transactions and roll them back for each test. Not working for me:

from django.test import TestCase

class MyTests(TestCase):
    def test_yair_a(self):
        print 'AAAAAAA'
        print Account.objects.all()
        Account.objects.create(account_id=1,account_name='1')
        print Account.objects.all()

    def test_yair_b(self):
        print 'BBBBBBBB'
        print Account.objects.all()
        Account.objects.create(account_id=2,account_name='2')
        print Account.objects.all()

Result (the interesting parts):

> py.test -s -v -k test_yair
AAAAAAA
[]
[<Account: 1>]
PASSED

BBBBBBBB
[<Account: 1>]
[<Account: 1>, <Account: 2>]
PASSED

Meaning no transaction was rolled back at the end of test_a.

Upvotes: 1

Views: 2766

Answers (1)

reformy
reformy

Reputation: 1027

The problem was that by default django creates a transaction on the default DB. I need a transaction on another DB. This is my settings file:

ROOT_PASS = 'ohnoyoudidnt'
ROOT_USER = 'root'
ROOT_HOST = 'localhost'

DATABASES = {
    'default': {
        'ENGINE': 'django.db.backends.mysql',
        'NAME': 'default_db',
        'USER': ROOT_USER,
        'PASSWORD': ROOT_PASS,
        'HOST': ROOT_HOST,
        'PORT': '3306',
    },
    'my_app': {
        'ENGINE': 'django.db.backends.mysql',
        'NAME': 'my_app_db',
        'USER': ROOT_USER,
        'PASSWORD': ROOT_PASS,
        'HOST': ROOT_HOST,
        'PORT': '3306'
    },

To solve this, one needs to use the transaction.atomic(using='my_app') block, and throw an error in it to rollback. Naturally I wanted to write a decorator for that, but using it with pytest was not as simple as I first thought, since pytest itself manipulates the method before my code (when using mocks and patches).

After spending a few hours, I got this decorator that works!

from decorator import decorator
import pytest
from django.db import transaction


class TestTempException(Exception):
    pass


def testing_w_transaction(using):
    def testing_w_transaction_inner(function):
        def func_wrapper(f, *args, **kwargs):
            try:
                with transaction.atomic(using=using):
                    f(*args, **kwargs)
                    raise TestTempException()
            except TestTempException:
                pass
        return decorator(func_wrapper, pytest.mark.django_db(function))
    return testing_w_transaction_inner

Usage:

@testing_w_transaction(using='my_app')
def test_my_method(mock_object, ...):
     # test test test

Upvotes: 1

Related Questions