TMP: better tests

src/olympia/amo/monitors.py

@@ -3,6 +3,7 @@
 import socket
 import traceback
 from urllib.parse import urljoin
+import MySQLdb
 
 from django.conf import settings
 
@@ -224,7 +225,6 @@ def signer():
 
 def olympia_database():
     """Check database connection by verifying the olympia database exists."""
-    import MySQLdb as mysql
 
     status = ''
 
@@ -243,14 +243,18 @@ def olympia_database():
     if db_info.get('PORT'):
         mysql_args['port'] = int(db_info.get('PORT'))
 
-    connection = mysql.connect(**mysql_args)
+    try:
+        connection = MySQLdb.connect(**mysql_args)
 
-    expected_db_name = db_info.get('NAME')
-    connection.query(f'SHOW DATABASES LIKE "{expected_db_name}"')
-    result = connection.store_result()
+        expected_db_name = db_info.get('NAME')
+        connection.query(f'SHOW DATABASES LIKE "{expected_db_name}"')
+        result = connection.store_result()
 
-    if result.num_rows() == 0:
-        status = f'Database {expected_db_name} does not exist'
+        if result.num_rows() == 0:
+            status = f'Database {expected_db_name} does not exist'
+            monitor_log.critical(status)
+    except Exception as e:
+        status = f'Failed to connect to database: {e}'
         monitor_log.critical(status)
 
     return status, None

src/olympia/amo/tests/test_monitor.py

@@ -191,3 +191,25 @@ def test_celery_worker_success(self):
             status, result = monitors.celery_worker()
             assert status == ''
             assert result is None
+
+    @patch('olympia.amo.monitors.MySQLdb.connect')
+    def test_olympia_db_unavailable(self, mock_connect):
+        mock_connect.side_effect = Exception('Database is unavailable')
+        status, result = monitors.olympia_database()
+        assert status == 'Failed to connect to database: Database is unavailable'
+        assert result is None
+
+    @patch('olympia.amo.monitors.MySQLdb.connect')
+    def test_olympia_db_database_does_not_exist(self, mock_connect):
+        # Create a mock connection object
+        mock_connection = MagicMock()
+        mock_connect.return_value = mock_connection
+
+        # Create a mock result object with num_rows returning 0
+        mock_result = MagicMock()
+        mock_result.num_rows.return_value = 0
+        mock_connection.store_result.return_value = mock_result
+
+        status, result = monitors.olympia_database()
+        assert status == 'Database test_olympia does not exist'
+        assert result is None

src/olympia/core/apps.py

@@ -125,18 +125,26 @@ def static_check(app_configs, **kwargs):
 def db_charset_check(app_configs, **kwargs):
     errors = []
 
-    with connection.cursor() as cursor:
-        cursor.execute("SHOW VARIABLES LIKE 'character_set_database';")
-        result = cursor.fetchone()
-        if result[1] != settings.DB_CHARSET:
-            errors.append(
-                Error(
-                    'Database charset invalid. '
-                    f'Expected {settings.DB_CHARSET}, '
-                    f'recieved {result[1]}',
-                    id='setup.E005',
+    try:
+        with connection.cursor() as cursor:
+            cursor.execute("SHOW VARIABLES LIKE 'character_set_database';")
+            result = cursor.fetchone()
+            if result[1] != settings.DB_CHARSET:
+                errors.append(
+                    Error(
+                        'Database charset invalid. '
+                        f'Expected {settings.DB_CHARSET}, '
+                        f'recieved {result[1]}',
+                        id='setup.E005',
+                    )
                 )
+    except Exception as e:
+        errors.append(
+            Error(
+                f'Failed to connect to database: {e}',
+                id='setup.E006',
             )
+        )
 
     return errors
 

src/olympia/core/tests/test_apps.py

@@ -36,6 +36,15 @@ def test_db_charset_check(self, mock_cursor):
         ):
             call_command('check')
 
+    @mock.patch('olympia.core.apps.connection.cursor')
+    def test_db_unavailable_check(self, mock_cursor):
+        mock_cursor.side_effect = Exception('Database is unavailable')
+        with self.assertRaisesMessage(
+            SystemCheckError,
+            'Failed to connect to database: Database is unavailable',
+        ):
+            call_command('check')
+
     def test_uwsgi_check(self):
         call_command('check')