Merge "Fix DB2 unit test to properly mock restart"

1 files changed, 7 insertions, 3 deletions

diff --git a/trove/tests/unittests/guestagent/test_db2_manager.py b/trove/tests/unittests/guestagent/test_db2_manager.py
index f1d8dac4..e965d58a 100644
--- a/trove/tests/unittests/guestagent/test_db2_manager.py
+++ b/trove/tests/unittests/guestagent/test_db2_manager.py
@@ -14,6 +14,7 @@
 
 import testtools
 from mock import MagicMock
+from mock import patch
 from testtools.matchers import Is, Equals, Not
 from trove.common.context import TroveContext
 from trove.common.instance import ServiceStatuses
@@ -132,9 +133,12 @@ class GuestAgentDB2ManagerTest(testtools.TestCase):
     def test_restart(self):
         mock_status = MagicMock()
         self.manager.appStatus = mock_status
-        db2_service.DB2App.restart = MagicMock(return_value=None)
-        self.manager.restart(self.context)
-        db2_service.DB2App.restart.assert_any_call()
+        with patch.object(db2_service.DB2App, 'restart',
+                          return_value=None) as restart_mock:
+            #invocation
+            self.manager.restart(self.context)
+            #verification/assertion
+            restart_mock.assert_any_call()
 
     def test_stop_db(self):
         mock_status = MagicMock()