From 48cf6242a8a9c57f6de7fe43fb3b2145aba39713 Mon Sep 17 00:00:00 2001
From: wudongze <wudongze@kanyun.com>
Date: Mon, 15 Jul 2024 16:52:25 +0800
Subject: [PATCH] db test 2 failure

---
 app/app/settings.py                         |  1 +
 app/core/management/__init__.py             |  0
 app/core/management/commands/__init__.py    |  0
 app/core/management/commands/wait_for_db.py | 23 +++++++++++++++++
 app/core/tests.py                           |  3 ---
 app/core/tests/__init__.py                  |  0
 app/core/tests/test_commands.py             | 28 +++++++++++++++++++++
 app/core/views.py                           |  3 ---
 8 files changed, 52 insertions(+), 6 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
 delete mode 100644 app/core/tests.py
 create mode 100644 app/core/tests/__init__.py
 create mode 100644 app/core/tests/test_commands.py
 delete mode 100644 app/core/views.py

diff --git a/app/app/settings.py b/app/app/settings.py
index 7f73f3d..ad58303 100644
--- a/app/app/settings.py
+++ b/app/app/settings.py
@@ -37,6 +37,7 @@
     'django.contrib.sessions',
     'django.contrib.messages',
     'django.contrib.staticfiles',
+    'core'
 ]
 
 MIDDLEWARE = [
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..fbbe338
--- /dev/null
+++ b/app/core/management/commands/wait_for_db.py
@@ -0,0 +1,23 @@
+"""
+Django command to wait for the db to be available.
+"""
+import time
+from django.core.management.base import BaseCommand
+from psycopg2 import OperationalError as Psycopg2OpError
+from django.db.utils import OperationalError
+
+
+class Command(BaseCommand):
+
+    def handle(self,*args,**options):
+        self.stdout.write('Waiting for database...')
+        db_up = False
+        while db_up is False:
+            try:
+                self.check(databases=['default'])
+                db_up = True
+            except(Psycopg2OpError, OperationalError):
+                self.stdout.write('Database unavailable, waiting 1 second...')
+                time.sleep(1)
+
+        self.stdout.write(self.style.SUCCESS('Database available!'))  
\ No newline at end of file
diff --git a/app/core/tests.py b/app/core/tests.py
deleted file mode 100644
index 7ce503c..0000000
--- a/app/core/tests.py
+++ /dev/null
@@ -1,3 +0,0 @@
-from django.test import TestCase
-
-# Create your tests here.
diff --git a/app/core/tests/__init__.py b/app/core/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/app/core/tests/test_commands.py b/app/core/tests/test_commands.py
new file mode 100644
index 0000000..5b1a2b1
--- /dev/null
+++ b/app/core/tests/test_commands.py
@@ -0,0 +1,28 @@
+"""
+ Test Django management commands.
+"""
+from unittest.mock import patch
+
+from psycopg2 import OperationalError as Psycopg2Error
+
+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):
+        """Test waiting for db when db is available"""
+        patched_check.return_value = True
+        call_command('wait_for_db')
+        patched_check.assert_called_once_with(database=['default'])
+
+    @patch('time.sleep')
+    def test_wait_for_db_delay(self, patched_sleep, patched_check):
+        """Test waiting for db when getting OperationalError"""
+        patched_check.side_effect = [Psycopg2Error] * 2 + \
+            [OperationalError] * 3 + [True]
+        call_command('wait_for_db')
+        self.assertEqual(patched_check.call_count, 6)
+        patched_check.assert_called_with(database=['default'])
\ No newline at end of file
diff --git a/app/core/views.py b/app/core/views.py
deleted file mode 100644
index 91ea44a..0000000
--- a/app/core/views.py
+++ /dev/null
@@ -1,3 +0,0 @@
-from django.shortcuts import render
-
-# Create your views here.