Create "querycount" parameter

Author: renanivo

pytest_django/plugin.py

@@ -83,6 +83,11 @@ def pytest_addoption(parser):
                   'Fail for invalid variables in templates.',
                   type='bool', default=False)
 
+    group._addoption('--querycount',
+                     action='store', type=int, default=None, metavar='N',
+                     help='Show top N tests with most queries '
+                          '(N=0 for all).')
+
 
 def _exists(path, ignore=EnvironmentError):
     try:
@@ -317,6 +322,59 @@ def pytest_runtest_setup(item):
         _disable_class_methods(cls)
 
 
+@pytest.hookimpl(hookwrapper=True)
+def pytest_runtest_call(item):
+    count_parameter = item.config.option.querycount
+    if count_parameter is None:
+        yield
+        return
+
+    from django.test.utils import CaptureQueriesContext
+    from django.db import connection
+
+    with CaptureQueriesContext(connection) as context:
+        yield
+
+    item.add_report_section('call', 'queries', context.captured_queries)
+
+
+def pytest_terminal_summary(terminalreporter, exitstatus):
+    count_parameter = terminalreporter.config.option.querycount
+    if count_parameter is None:
+        return
+
+    if count_parameter:
+        header = 'top {} tests with most queries'.format(count_parameter)
+        reports_slice = slice(None, count_parameter)
+    else:
+        header = 'top tests with most queries'
+        reports_slice = slice(None, None)
+
+    terminalreporter.write_sep('=', header)
+
+    def get_query_count(report):
+        sections = dict(report.sections)
+        return len(sections.get('Captured queries call', []))
+
+    reports = (
+        terminalreporter.stats.get('failed', []) +
+        terminalreporter.stats.get('passed', [])
+    )
+
+    reports.sort(key=get_query_count)
+    reports.reverse()
+
+    for report in reports[reports_slice]:
+        count = get_query_count(report)
+        nodeid = report.nodeid.replace("::()::", "::")
+
+        terminalreporter.write_line('{count: <4} {when: <8} {nodeid}'.format(
+            count=count,
+            when=report.when,
+            nodeid=nodeid
+        ))
+
+
 @pytest.fixture(autouse=True, scope='session')
 def django_test_environment(request):
     """

tests/test_report.py

@@ -0,0 +1,90 @@
+class TestQueryCount(object):
+    """Test report generated by --querycount parameter"""
+
+    def test_querycount_report_header(self, django_testdir):
+        django_testdir.create_test_module('''
+            def test_zero_queries():
+                pass
+        ''')
+
+        result = django_testdir.runpytest_subprocess('--querycount=5')
+        result.stdout.fnmatch_lines([
+            '*== top 5 tests with most queries ==*'
+        ])
+
+    def test_header_not_set_without_parameter(self, django_testdir):
+        django_testdir.create_test_module('''
+            def test_zero_queries():
+                pass
+        ''')
+
+        result = django_testdir.runpytest_subprocess()
+        assert 'tests with most queries' not in result.stdout.str()
+
+    def test_querycount_report_lines(self, django_testdir):
+        django_testdir.create_test_module('''
+            import pytest
+            from django.db import connection
+
+            @pytest.mark.django_db
+            def test_one_query():
+                with connection.cursor() as cursor:
+                    cursor.execute('SELECT 1')
+
+                assert True
+
+            @pytest.mark.django_db
+            def test_two_queries():
+                with connection.cursor() as cursor:
+                    cursor.execute('SELECT 1')
+                    cursor.execute('SELECT 1')
+
+                assert True
+
+            @pytest.mark.django_db
+            def test_failed_one_query():
+                with connection.cursor() as cursor:
+                    cursor.execute('SELECT 1')
+
+                assert False
+
+            def test_zero_queries():
+                assert True
+        ''')
+
+        result = django_testdir.runpytest_subprocess('--querycount=3')
+        lines = result.stdout.get_lines_after(
+            '*top 3 tests with most queries*'
+        )
+        assert 'test_two_queries' in lines[0]
+        assert 'test_one_query' in lines[1]
+        assert 'test_failed' in lines[2]
+        assert 'test_zero_queries' not in result.stdout.str()
+
+    def test_report_all_lines_on_querycount_zero(self, django_testdir):
+        django_testdir.create_test_module('''
+            import pytest
+            from django.db import connection
+
+            @pytest.mark.django_db
+            def test_one_query():
+                with connection.cursor() as cursor:
+                    cursor.execute('SELECT 1')
+
+                assert True
+
+            @pytest.mark.django_db
+            def test_two_queries():
+                with connection.cursor() as cursor:
+                    cursor.execute('SELECT 1')
+                    cursor.execute('SELECT 1')
+
+                assert True
+        ''')
+
+        result = django_testdir.runpytest_subprocess('--querycount=0')
+        lines = result.stdout.get_lines_after(
+            '*top tests with most queries*'
+        )
+        assert 'test_two_queries' in lines[0]
+        assert 'test_one_query' in lines[1]