db test 2 failure

app/app/settings.py

@@ -37,6 +37,7 @@
     'django.contrib.sessions',
     'django.contrib.messages',
     'django.contrib.staticfiles',
+    'core'
 ]
 
 MIDDLEWARE = [

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!'))  

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'])