Feat: wait_for_db 커맨드 테스트 코드 추가

core/management/commands/wait_for_db.py

@@ -10,14 +10,15 @@ class Command(BaseCommand):
 
     def handle(self, *args, **options):
         for i in range(10):
-            self.stdout.write(f'Try {i}: Waiting for database...')
+            self.stdout.write(f'Try {i+1}: Waiting for database...')
             try:
                 connection = connections['default']
                 if connection:
                     self.stdout.write(self.style.SUCCESS('PostgreSQL Database available!'))
                     return
             except (Psycopg2OpError, OperationalError):
                 if i == 9:
-                    self.stdout.write(self.style.FAIL(f"PostgreSQL Connection Failed after {i+1} attempts"))
-                self.stdout.write(self.style.FAIL("PostgreSQL Connection Failed! Connection Retrying..."))
-                time.sleep(1)
+                    self.stdout.write(self.style.ERROR(f"PostgreSQL Connection Failed after {i+1} attempts"))
+                else:
+                    self.stdout.write(self.style.ERROR("PostgreSQL Connection Failed! Connection Retrying..."))
+                time.sleep(1)

core/tests.py

@@ -1,3 +1,35 @@
-from django.test import TestCase
+from django.test import SimpleTestCase
+from django.db.utils import OperationalError
+from unittest.mock import patch
+from django.core.management import call_command
+from psycopg2 import OperationalError as Psycopg2OpError
 
-# Create your tests here.
+
+@patch('django.db.utils.ConnectionHandler.__getitem__')
+class CommandTests(SimpleTestCase):
+    # db연결에 한번에 성공했을 때
+    def test_wait_for_db_ready(self, patched_getitem):
+        patched_getitem.return_value = True
+
+        call_command('wait_for_db')
+
+        self.assertEqual(patched_getitem.call_count, 1)
+
+    # db연결이 지연될 때
+    @patch('time.sleep')
+    def test_wait_for_db_delay(self, patched_sleep, patched_getitem):
+        patched_getitem.side_effect = [Psycopg2OpError] \
+            + [OperationalError]*5 + [True]
+
+        call_command('wait_for_db')
+
+        self.assertEqual(patched_getitem.call_count, 7)
+
+    @patch('time.sleep')
+    def test_wait_for_db_failed(self, patched_sleep, patched_getitem):
+        patched_getitem.side_effect = [Psycopg2OpError] \
+                                      + [OperationalError] * 10
+
+        call_command('wait_for_db')
+
+        self.assertEqual(patched_getitem.call_count, 10)