Python nplusone inside tests

test_nplusone.py

from functools import wraps
from django.test import TestCase

def raise_nplusone_decorator(cls):
    original_setup_class = getattr(cls, 'setUpClass')
    original_teardown_class = getattr(cls, 'tearDownClass')

    @classmethod
    def new_setup_class(cls):
        # Perform setup steps before the test class
        with Profiler():
            original_setup_class()

    @classmethod
    def new_teardown_class(cls):
        # Perform teardown steps after the test class
        original_teardown_class()

    cls.setUpClass = new_setup_class
    cls.tearDownClass = new_teardown_class
    return cls

@raise_nplusone_decorator
class MyTestCase(TestCase):
    def test_my_functionality(self):
        # Your test logic goes here
        self.assertEqual(2 + 2, 4)