[mariosplivalo, r=timkuhlman] Code cleanup, removing obsolete references to unused libs

2 files changed, 25 insertions, 21 deletions

diff --git a/hooks/hooks.py b/hooks/hooks.py
index d54919c..30fc460 100755
--- a/hooks/hooks.py
+++ b/hooks/hooks.py
@@ -70,17 +70,11 @@ from charmhelpers.contrib.hahelpers.cluster import (
 )
 
 try:
- try:
- from pymongo import Connection
- except ImportError:
- from pymongo import MongoClient as Connection
+ from pymongo import MongoClient
 from pymongo.errors import OperationFailure
 except ImportError:
 apt_install("python-pymongo", fatal=True)
- try:
- from pymongo import Connection
- except ImportError:
- from pymongo import MongoClient as Connection
+ from pymongo import MongoClient
 from pymongo.errors import OperationFailure
 
 from charmhelpers.contrib.charmsupport.nrpe import NRPE
@@ -491,7 +485,7 @@ def init_replset():
 
 retVal = mongo_client('localhost', 'rs.initiate(%s)' % init)
 time.sleep(1) # give mongod some time to become primary
- c = Connection('localhost')
+ c = MongoClient('localhost')
 while True:
 try:
 r = run_admin_command(c, 'replSetGetStatus')
@@ -1165,7 +1159,7 @@ def replica_set_relation_joined():
 
 
 def am_i_primary():
- c = Connection('localhost')
+ c = MongoClient('localhost')
 for i in xrange(10):
 try:
 r = run_admin_command(c, 'replSetGetStatus')
@@ -1178,7 +1172,10 @@ def am_i_primary():
 # replSet initialization in progress
 continue
 elif 'EMPTYCONFIG' in str(e):
- # replication not
+ # replication not initialized yet
+ return False
+ elif 'not running with --replSet' in str(e):
+ # replicaset not configured 
 return False
 else:
 raise
@@ -1263,7 +1260,7 @@ def replica_set_relation_broken():
 mongo_client('localhost', 'rs.stepDown()')
 time.sleep(15) # give some time to for re-election to happen
 
- c = Connection('localhost')
+ c = MongoClient('localhost')
 r = c.admin.command('isMaster')
 
 try:
diff --git a/unit_tests/test_hooks.py b/unit_tests/test_hooks.py
index c95e480..603b149 100644
--- a/unit_tests/test_hooks.py
+++ b/unit_tests/test_hooks.py
@@ -1,12 +1,12 @@
 from mock import patch, call
 
-import hooks
-
 from test_utils import CharmTestCase
 from test_utils import mock_open
 from pymongo.errors import OperationFailure
 from subprocess import CalledProcessError
 
+import hooks
+
 import tempfile
 import os
 
@@ -74,7 +74,7 @@ class MongoHooksTest(CharmTestCase):
 self.relation_set.assert_called_with('fake-relation-id', exp_rel_vals)
 
 @patch.object(hooks, 'run_admin_command')
- @patch.object(hooks, 'Connection')
+ @patch.object(hooks, 'MongoClient')
 @patch.object(hooks, 'config')
 @patch.object(hooks, 'mongo_client')
 @patch('time.sleep')
@@ -161,7 +161,7 @@ class MongoHooksTest(CharmTestCase):
 fatal=True)
 
 @patch.object(hooks, 'run_admin_command')
- @patch.object(hooks, 'Connection')
+ @patch.object(hooks, 'MongoClient')
 @patch('time.sleep')
 def test_am_i_primary(self, mock_sleep, mock_mongo_client,
 mock_run_admin_cmd):
@@ -174,7 +174,7 @@ class MongoHooksTest(CharmTestCase):
 self.assertEqual(exp, rv)
 
 @patch.object(hooks, 'run_admin_command')
- @patch.object(hooks, 'Connection')
+ @patch.object(hooks, 'MongoClient')
 @patch('time.sleep')
 def test_am_i_primary_too_many_attempts(self, mock_sleep,
 mock_mongo_client,
@@ -191,13 +191,20 @@ class MongoHooksTest(CharmTestCase):
 pass
 
 @patch.object(hooks, 'run_admin_command')
- @patch.object(hooks, 'Connection')
+ @patch.object(hooks, 'MongoClient')
 @patch('time.sleep')
 def test_am_i_primary_operation_failures(self, mock_sleep,
 mock_mongo_client,
 mock_run_admin_cmd):
- mock_run_admin_cmd.side_effect = OperationFailure('EMPTYCONFIG')
+ msg = 'EMPTYCONFIG'
+ mock_run_admin_cmd.side_effect = OperationFailure(msg)
+ rv = hooks.am_i_primary()
+ self.assertTrue(mock_run_admin_cmd.called)
+ self.assertFalse(rv)
 
+ mock_run_admin_cmd.reset_mock()
+ msg = 'not running with --replSet'
+ mock_run_admin_cmd.side_effect = OperationFailure(msg)
 rv = hooks.am_i_primary()
 self.assertTrue(mock_run_admin_cmd.called)
 self.assertFalse(rv)
@@ -345,13 +352,13 @@ class MongoHooksTest(CharmTestCase):
 mock_leave_replset.assert_has_calls([call1])
 
 @patch('time.sleep')
- @patch.object(hooks, 'Connection')
+ @patch.object(hooks, 'MongoClient')
 @patch.object(hooks, 'unit_get')
 @patch.object(hooks, 'leave_replset')
 @patch.object(hooks, 'am_i_primary')
 def test_replica_set_relation_broken(self, mock_am_i_primary,
 mock_leave_replset, mock_unit_get,
- mock_Connection, mock_sleep):
+ mock_MongoClient, mock_sleep):
 
 mock_am_i_primary.return_value = False
 hooks.replica_set_relation_broken()