File: test_debugsqlshell.py

package info (click to toggle)
python-django-debug-toolbar 1%3A6.0.0-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 2,052 kB
  • sloc: python: 7,555; javascript: 636; makefile: 67; sh: 16
file content (36 lines) | stat: -rw-r--r-- 1,244 bytes parent folder | download | duplicates (3)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
import io
import sys

from django.contrib.auth.models import User
from django.core import management
from django.db import connection
from django.test import TestCase
from django.test.utils import override_settings

if connection.vendor == "postgresql":
    from django.db.backends.postgresql import base as base_module
else:
    from django.db.backends import utils as base_module


@override_settings(DEBUG=True)
class DebugSQLShellTestCase(TestCase):
    def setUp(self):
        self.original_wrapper = base_module.CursorDebugWrapper
        # Since debugsqlshell monkey-patches django.db.backends.utils, we can
        # test it simply by loading it, without executing it. But we have to
        # undo the monkey-patch on exit.
        command_name = "debugsqlshell"
        app_name = management.get_commands()[command_name]
        management.load_command_class(app_name, command_name)

    def tearDown(self):
        base_module.CursorDebugWrapper = self.original_wrapper

    def test_command(self):
        original_stdout, sys.stdout = sys.stdout, io.StringIO()
        try:
            User.objects.count()
            self.assertIn("SELECT COUNT", sys.stdout.getvalue())
        finally:
            sys.stdout = original_stdout