Gracefully disable Merlin support if MySQLdb isn't availabe

marineam · Sep 27, 2011 · ae8c12c · ae8c12c
1 parent 4a8cb6d
commit ae8c12c
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 14 deletions.
diff --git a/python/nagcat/main.py b/python/nagcat/main.py
@@ -132,6 +132,9 @@ def parse_options():
     # Flag to use merlin or not
     options.merlin = any_merlin_settings and all_merlin_settings
 
+    if options.merlin and not merlin.available():
+        err.append("Merlin support is not available. Missing MySQLdb?")
+
     if any_merlin_settings and not all_merlin_settings:
         err.append("All merlin database settings must be included")
 

diff --git a/python/nagcat/merlin.py b/python/nagcat/merlin.py
@@ -14,17 +14,27 @@
 
 import os
 import errno
-import MySQLdb
 import time
 
+try:
+    import MySQLdb
+except ImportError:
+    MySQLdb = None
+
 from coil.errors import CoilError
 from nagcat import errors, log, nagios_api
 from nagcat import nagios_objects, scheduler, merlintest, nagios, simple
 
+
+def available():
+    """Returns False if merlin support is disabled"""
+    return MySQLdb is not None
+
 class NagcatMerlin(nagios.NagcatNagios):
     """NagcatNagios scheduler that load balances using merlin."""
 
     def __init__(self, config, nagios_cfg, merlin_db_info={}, **kwargs):
+        assert available()
         nagios.NagcatNagios.__init__(self, config, nagios_cfg, **kwargs)
         self._test_index = 0
         self._merlin_db_info = merlin_db_info

diff --git a/python/nagcat/unittests/test_merlin.py b/python/nagcat/unittests/test_merlin.py
@@ -17,23 +17,28 @@
 from coil.struct import Struct
 import os
 import warnings
-import MySQLdb
+
+try:
+    import MySQLdb
+except ImportError:
+    skip = "MySQLdb is not available"
 
 class TestNagcatMerlinCase(unittest.TestCase):
 
+    _merlin_db_info = {
+        'merlin_db_user': os.environ.get('MYSQL_USER', None),
+        'merlin_db_host': os.environ.get('MYSQL_HOST', None),
+        'merlin_db_pass': os.environ.get('MYSQL_PASS', None),
+        'merlin_db_name': os.environ.get('MYSQL_NAME', None),
+        'merlin_db_port': os.environ.get('MYSQL_PORT', 3306),
+    }
+
+    if not all((_merlin_db_info['merlin_db_user'],
+                _merlin_db_info['merlin_db_host'],
+                _merlin_db_info['merlin_db_name'])):
+        skip = "Missing MySQL credentials"
+
     def setUp(self):
-        self._merlin_db_info = {
-            "merlin_db_user": os.environ.get('MYSQL_USER', None),
-            "merlin_db_host": os.environ.get('MYSQL_HOST', None),
-            "merlin_db_pass": os.environ.get('MYSQL_PASS', None),
-            "merlin_db_name": os.environ.get('MYSQL_NAME', None),
-            "merlin_db_port": os.environ.get('MYSQL_PORT', 3306),
-        }
-        self._should_skip = not all([self._merlin_db_info['merlin_db_user'],
-            self._merlin_db_info['merlin_db_host'],
-            self._merlin_db_info['merlin_db_name']])
-        if self._should_skip:
-            raise unittest.SkipTest("Not enough database information")
         db = MySQLdb.connect(
             user=self._merlin_db_info['merlin_db_user'],
             host=self._merlin_db_info['merlin_db_host'],