Skip to content

Commit

Permalink
fixtures: ensure db blocker is always restored
Browse files Browse the repository at this point in the history
As the code is written currently, it's possible the `restore` is not
called in case of exception. Use the context manager to ensure that it
does.

This is not to fix an observed issue, just a fix from reviewing the
code.
  • Loading branch information
bluetech committed Sep 28, 2024
1 parent 113b578 commit 263ca6d
Showing 1 changed file with 52 additions and 55 deletions.
107 changes: 52 additions & 55 deletions pytest_django/fixtures.py
Original file line number Diff line number Diff line change
Expand Up @@ -197,66 +197,63 @@ def _django_db_helper(
"django_db_serialized_rollback" in request.fixturenames
)

django_db_blocker.unblock()

import django.db
import django.test

if transactional:
test_case_class = django.test.TransactionTestCase
else:
test_case_class = django.test.TestCase

_reset_sequences = reset_sequences
_serialized_rollback = serialized_rollback
_databases = databases
_available_apps = available_apps

class PytestDjangoTestCase(test_case_class): # type: ignore[misc,valid-type]
reset_sequences = _reset_sequences
serialized_rollback = _serialized_rollback
if _databases is not None:
databases = _databases
if _available_apps is not None:
available_apps = _available_apps

# For non-transactional tests, skip executing `django.test.TestCase`'s
# `setUpClass`/`tearDownClass`, only execute the super class ones.
#
# `TestCase`'s class setup manages the `setUpTestData`/class-level
# transaction functionality. We don't use it; instead we (will) offer
# our own alternatives. So it only adds overhead, and does some things
# which conflict with our (planned) functionality, particularly, it
# closes all database connections in `tearDownClass` which inhibits
# wrapping tests in higher-scoped transactions.
#
# It's possible a new version of Django will add some unrelated
# functionality to these methods, in which case skipping them completely
# would not be desirable. Let's cross that bridge when we get there...
if not transactional:

@classmethod
def setUpClass(cls) -> None:
super(django.test.TestCase, cls).setUpClass()

@classmethod
def tearDownClass(cls) -> None:
super(django.test.TestCase, cls).tearDownClass()

PytestDjangoTestCase.setUpClass()

test_case = PytestDjangoTestCase(methodName="__init__")
test_case._pre_setup()
with django_db_blocker.unblock():
import django.db
import django.test

yield
if transactional:
test_case_class = django.test.TransactionTestCase
else:
test_case_class = django.test.TestCase

_reset_sequences = reset_sequences
_serialized_rollback = serialized_rollback
_databases = databases
_available_apps = available_apps

class PytestDjangoTestCase(test_case_class): # type: ignore[misc,valid-type]
reset_sequences = _reset_sequences
serialized_rollback = _serialized_rollback
if _databases is not None:
databases = _databases
if _available_apps is not None:
available_apps = _available_apps

# For non-transactional tests, skip executing `django.test.TestCase`'s
# `setUpClass`/`tearDownClass`, only execute the super class ones.
#
# `TestCase`'s class setup manages the `setUpTestData`/class-level
# transaction functionality. We don't use it; instead we (will) offer
# our own alternatives. So it only adds overhead, and does some things
# which conflict with our (planned) functionality, particularly, it
# closes all database connections in `tearDownClass` which inhibits
# wrapping tests in higher-scoped transactions.
#
# It's possible a new version of Django will add some unrelated
# functionality to these methods, in which case skipping them completely
# would not be desirable. Let's cross that bridge when we get there...
if not transactional:

@classmethod
def setUpClass(cls) -> None:
super(django.test.TestCase, cls).setUpClass()

@classmethod
def tearDownClass(cls) -> None:
super(django.test.TestCase, cls).tearDownClass()

PytestDjangoTestCase.setUpClass()

test_case = PytestDjangoTestCase(methodName="__init__")
test_case._pre_setup()

test_case._post_teardown()
yield

PytestDjangoTestCase.tearDownClass()
test_case._post_teardown()

PytestDjangoTestCase.doClassCleanups()
PytestDjangoTestCase.tearDownClass()

django_db_blocker.restore()
PytestDjangoTestCase.doClassCleanups()


def validate_django_db(marker: pytest.Mark) -> _DjangoDb:
Expand Down

0 comments on commit 263ca6d

Please sign in to comment.