/
test_debugsqlshell.py
34 lines (29 loc) · 1.29 KB
/
test_debugsqlshell.py
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
import io
import sys
from django.contrib.auth.models import User
from django.core import management
from django.db.backends import utils as db_backends_utils
from django.db.backends.postgresql import base as psql_base
from django.test import TestCase
from django.test.utils import override_settings
@override_settings(DEBUG=True)
class DebugSQLShellTestCase(TestCase):
def setUp(self):
self.original_utils_wrapper = db_backends_utils.CursorDebugWrapper
self.original_psql_wrapper = psql_base.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):
db_backends_utils.CursorDebugWrapper = self.original_utils_wrapper
psql_base.CursorDebugWrapper = self.original_psql_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