From fa3e5f391ed570f2772746d2f44129f34c2c6d83 Mon Sep 17 00:00:00 2001
From: Penny <enderman542@gmail.com>
Date: Sat, 2 Nov 2024 13:43:04 +0900
Subject: [PATCH] add wait_for_db command

---
 app/app/settings.py                         |  2 +-
 app/core/admin.py                           |  2 +-
 app/core/management/__init__.py             |  0
 app/core/management/commands/__init__.py    |  0
 app/core/management/commands/wait_for_db.py | 22 +++++++++++++++
 app/core/models.py                          |  2 +-
 app/core/tests/test_commands.py             | 31 +++++++++++++++++++++
 docker-compose.yml                          |  2 +-
 8 files changed, 57 insertions(+), 4 deletions(-)
 create mode 100644 app/core/management/__init__.py
 create mode 100644 app/core/management/commands/__init__.py
 create mode 100644 app/core/management/commands/wait_for_db.py
 create mode 100644 app/core/tests/test_commands.py

diff --git a/app/app/settings.py b/app/app/settings.py
index 5b9ca25..56cadcd 100644
--- a/app/app/settings.py
+++ b/app/app/settings.py
@@ -81,7 +81,7 @@
         'HOST': os.environ.get('DB_HOST'),
         'NAME': os.environ.get('DB_NAME'),
         'USER': os.environ.get('DB_USER'),
-        'PASSWORD': os.environ.get('DB_PASS'),
+        'PASSWORD': os.environ.get('DB_PASSWORD'),
     }
 }
 
diff --git a/app/core/admin.py b/app/core/admin.py
index 8c38f3f..a4e11e9 100644
--- a/app/core/admin.py
+++ b/app/core/admin.py
@@ -1,3 +1,3 @@
-from django.contrib import admin
+from django.contrib import admin # noqa
 
 # Register your models here.
diff --git a/app/core/management/__init__.py b/app/core/management/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/app/core/management/commands/__init__.py b/app/core/management/commands/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/app/core/management/commands/wait_for_db.py b/app/core/management/commands/wait_for_db.py
new file mode 100644
index 0000000..414646b
--- /dev/null
+++ b/app/core/management/commands/wait_for_db.py
@@ -0,0 +1,22 @@
+import time
+
+from psycopg2 import OperationalError as Psycopg2OpError
+
+from django.db.utils import OperationalError
+from django.core.management.base import BaseCommand
+
+
+class Command(BaseCommand):
+
+    def handle(self, *args, **options):
+        self.stdout.write('Waiting for db...')
+        db_up = False
+        while db_up is False:
+            try:
+                self.check(databases=['default'])
+                db_up = True
+            except (Psycopg2OpError, OperationalError):
+                self.stdout.write('db unavailable, waiting 1 second...')
+                time.sleep(1)
+
+        self.stdout.write(self.style.SUCCESS('db available!'))
diff --git a/app/core/models.py b/app/core/models.py
index 71a8362..56b48a5 100644
--- a/app/core/models.py
+++ b/app/core/models.py
@@ -1,3 +1,3 @@
-from django.db import models
+from django.db import models # noqa
 
 # Create your models here.
diff --git a/app/core/tests/test_commands.py b/app/core/tests/test_commands.py
new file mode 100644
index 0000000..7d3a097
--- /dev/null
+++ b/app/core/tests/test_commands.py
@@ -0,0 +1,31 @@
+"""
+Test custom Django management commands.
+"""
+from unittest.mock import patch
+
+from psycopg2 import OperationalError as Psycopg2OpError
+
+from django.core.management import call_command
+from django.db.utils import OperationalError
+from django.test import SimpleTestCase
+
+
+@patch('core.management.commands.wait_for_db.Command.check')
+class CommandTests(SimpleTestCase):
+
+    def test_wait_for_db_ready(self, patched_check):
+        patched_check.return_value = True
+
+        call_command('wait_for_db')
+
+        patched_check.assert_called_once_with(databases=['default'])
+
+    @patch('time.sleep')
+    def test_wait_for_db_delay(self, patched_sleep, patched_check):
+        patched_check.side_effect = [Psycopg2OpError] * 2 + \
+            [OperationalError] * 3 + [True]
+
+        call_command('wait_for_db')
+
+        self.assertEqual(patched_check.call_count, 6)
+        patched_check.assert_called_with(databases=['default'])
diff --git a/docker-compose.yml b/docker-compose.yml
index 54ddf70..c121c35 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -16,7 +16,7 @@ services:
       - DB_HOST=db
       - DB_NAME=devdb
       - DB_USER=devuser
-      - DB_PASS=Password1
+      - DB_PASSWORD=Password1
     depends_on:
       - db