Repository: openwisp/openwisp-utils (language: python)

Issue:
[test] AssertNumQueries as subtest
We need a way to run `assertNumQueries` as a subtest so it doesn't fail the whole test if the assertion fails, but lets the test finish.


Discussion / hints:

Fix patch:
diff --git a/README.rst b/README.rst
index b2ba8324..00a257ae 100644
--- a/README.rst
+++ b/README.rst
@@ -804,6 +804,30 @@ Example usage:
         self.assertIn('expected stdout', captured_output.getvalue())
         self.assertIn('expected stderr', captured_error.getvalue())
 
+``openwisp_utils.tests.AssertNumQueriesSubTestMixin``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+This mixin overrides the
+`assertNumQueries <https://docs.djangoproject.com/en/dev/topics/testing/tools/#django.test.TransactionTestCase.assertNumQueries>`_
+assertion from the django test case to run in a ``subTest`` so that the
+query check does not block the whole test if it fails.
+
+Example usage:
+
+.. code-block:: python
+
+    from django.test import TestCase
+    from openwisp_utils.tests import AssertNumQueriesSubTestMixin
+
+
+    class MyTest(AssertNumQueriesSubTestMixin, TestCase):
+        def my_test(self):
+            with self.assertNumQueries(2):
+                MyModel.objects.count()
+
+            # the assertion above will fail but this line will be executed
+            print('This will be printed anyway.')
+
 Quality Assurance Checks
 ------------------------
 


Test patch:
diff --git a/openwisp_utils/tests.py b/openwisp_utils/tests.py
index f919a2c2..41826295 100644
--- a/openwisp_utils/tests.py
+++ b/openwisp_utils/tests.py
@@ -5,7 +5,9 @@
 from unittest import TextTestResult, mock
 
 from django.conf import settings
+from django.db import DEFAULT_DB_ALIAS, connections
 from django.test.runner import DiscoverRunner
+from django.test.utils import CaptureQueriesContext
 
 from .utils import print_color
 
@@ -125,3 +127,46 @@ def __init__(self, stdout=None, stderr=None, **kwargs):
         self.kwargs = kwargs
         self.stdout = stdout or io.StringIO()
         self.stderr = stderr or io.StringIO()
+
+
+class _AssertNumQueriesContextSubTest(CaptureQueriesContext):
+    """
+    Needed to execute assertNumQueries in a subTest
+    """
+
+    def __init__(self, test_case, num, connection):
+        self.test_case = test_case
+        self.num = num
+        super().__init__(connection)
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        super().__exit__(exc_type, exc_value, traceback)
+        if exc_type is not None:
+            return
+        executed = len(self)
+        with self.test_case.subTest(f'Expecting {self.num} SQL queries'):
+            self.test_case.assertEqual(
+                executed,
+                self.num,
+                "%d queries executed, %d expected\nCaptured queries were:\n%s"
+                % (
+                    executed,
+                    self.num,
+                    '\n'.join(
+                        '%d. %s' % (i, query['sql'])
+                        for i, query in enumerate(self.captured_queries, start=1)
+                    ),
+                ),
+            )
+
+
+class AssertNumQueriesSubTestMixin:
+    def assertNumQueries(self, num, func=None, *args, using=DEFAULT_DB_ALIAS, **kwargs):
+        conn = connections[using]
+
+        context = _AssertNumQueriesContextSubTest(self, num, conn)
+        if func is None:
+            return context
+
+        with context:
+            func(*args, **kwargs)
diff --git a/tests/test_project/tests/test_test_utils.py b/tests/test_project/tests/test_test_utils.py
index eadcf1ec..aa29f5cc 100644
--- a/tests/test_project/tests/test_test_utils.py
+++ b/tests/test_project/tests/test_test_utils.py
@@ -1,8 +1,10 @@
 import sys
+import unittest
 
 from django.dispatch import Signal
 from django.test import TestCase, override_settings
 from openwisp_utils.tests import (
+    AssertNumQueriesSubTestMixin,
     TimeLoggingTestRunner,
     capture_any_output,
     capture_stderr,
@@ -11,6 +13,8 @@
 )
 from openwisp_utils.utils import deep_merge_dicts, print_color
 
+from ..models import Shelf
+
 status_signal = Signal(providing_args=['status'])
 
 
@@ -82,3 +86,13 @@ def test_print_color(self, captured_output):
     def test_capture_stderr(self, captured_error):
         print('Testing capture_stderr', file=sys.stderr, end='')
         self.assertEqual(captured_error.getvalue(), 'Testing capture_stderr')
+
+
+class TestAssertNumQueriesSubTest(AssertNumQueriesSubTestMixin, TestCase):
+    def test_assert_num_queries(self):
+        with unittest.mock.patch.object(
+            self, 'subTest', wraps=self.subTest
+        ) as patched_subtest:
+            with self.assertNumQueries(1):
+                Shelf.objects.count()
+            patched_subtest.assert_called_once()