From ed64b96ede0dbca5d675e1011bd979c7c3541e37 Mon Sep 17 00:00:00 2001 From: Brett Date: Tue, 8 Feb 2022 08:19:54 -0500 Subject: [PATCH 1/2] obtain cursor to test db ready --- app/core/management/commands/wait_for_db.py | 3 +++ app/core/tests/test_commands.py | 17 +++++++++-------- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/app/core/management/commands/wait_for_db.py b/app/core/management/commands/wait_for_db.py index b6bdcde..c37f1e8 100644 --- a/app/core/management/commands/wait_for_db.py +++ b/app/core/management/commands/wait_for_db.py @@ -14,8 +14,11 @@ def handle(self, *args, **options): while not db_conn: try: db_conn = connections['default'] + self.stdout.write('Obtain cursor and verify database...') + db_conn.cursor() except OperationalError: self.stdout.write('Database unavailable, waiting 1 second...') time.sleep(1) + db_conn = None self.stdout.write(self.style.SUCCESS('Database available!')) diff --git a/app/core/tests/test_commands.py b/app/core/tests/test_commands.py index 1ccd0d9..89c1ab5 100644 --- a/app/core/tests/test_commands.py +++ b/app/core/tests/test_commands.py @@ -1,4 +1,4 @@ -from unittest.mock import patch +from unittest.mock import patch, MagicMock from django.core.management import call_command from django.db.utils import OperationalError @@ -9,15 +9,16 @@ class CommandTests(TestCase): def test_wait_for_db_ready(self): """Test waiting for db when db is available""" - with patch('django.db.utils.ConnectionHandler.__getitem__') as gi: - gi.return_value = True + conn = MagicMock(return_value=None) + with patch('django.db.utils.ConnectionHandler.__getitem__', return_value=conn): call_command('wait_for_db') - self.assertEqual(gi.call_count, 1) + conn.cursor.assert_called_once() - @patch('time.sleep', return_value=True) + @patch('time.sleep', return_value=None) def test_wait_for_db(self, ts): """Test waiting for db""" - with patch('django.db.utils.ConnectionHandler.__getitem__') as gi: - gi.side_effect = [OperationalError] * 5 + [True] + conn = MagicMock(return_value=None) + with patch('django.db.utils.ConnectionHandler.__getitem__', return_value=conn): + conn.cursor.side_effect = [OperationalError] * 5 + [True] call_command('wait_for_db') - self.assertEqual(gi.call_count, 6) + self.assertEqual(conn.cursor.call_count, 6) From 67fef85fef4ff4b84672f7c8fbac57149112b793 Mon Sep 17 00:00:00 2001 From: Brett Date: Tue, 8 Feb 2022 08:25:39 -0500 Subject: [PATCH 2/2] loop on cursor --- app/core/management/commands/wait_for_db.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/app/core/management/commands/wait_for_db.py b/app/core/management/commands/wait_for_db.py index c37f1e8..cc5f1aa 100644 --- a/app/core/management/commands/wait_for_db.py +++ b/app/core/management/commands/wait_for_db.py @@ -10,15 +10,14 @@ class Command(BaseCommand): def handle(self, *args, **options): self.stdout.write('Waiting for database...') - db_conn = None - while not db_conn: + cursor = None + while not cursor: try: db_conn = connections['default'] self.stdout.write('Obtain cursor and verify database...') - db_conn.cursor() + cursor = db_conn.cursor() except OperationalError: self.stdout.write('Database unavailable, waiting 1 second...') time.sleep(1) - db_conn = None self.stdout.write(self.style.SUCCESS('Database available!'))