1 files changed, 18 insertions, 6 deletions
diff --git a/packages/python-google-compute-engine/google_compute_engine/distro_lib/el_6/tests/utils_test.py b/packages/python-google-compute-engine/google_compute_engine/distro_lib/el_6/tests/utils_test.py
index f47f5c7..1dcbc60 100644
--- a/packages/python-google-compute-engine/google_compute_engine/distro_lib/el_6/tests/utils_test.py
+++ b/packages/python-google-compute-engine/google_compute_engine/distro_lib/el_6/tests/utils_test.py
@@ -30,15 +30,27 @@ class UtilsTest(unittest.TestCase):
     pass
 
   @mock.patch('google_compute_engine.distro_lib.helpers.CallDhclientIpv6')
-  def testEnableIpv6(self, mock_call):
+  @mock.patch('google_compute_engine.distro_lib.helpers.CallEnableRouteAdvertisements')
+  def testEnableIpv6(self, mock_call_enable_ra, mock_call_dhclient):
     mocks = mock.Mock()
-    mocks.attach_mock(mock_call, 'call')
+    mocks.attach_mock(mock_call_dhclient, 'dhclient')
+    mocks.attach_mock(mock_call_enable_ra, 'enable_ra')
 
-    utils.Utils.EnableIpv6(
-        self.mock_setup, ['A', 'B'], self.mock_logger,
-        dhclient_script='test_script')
+    utils.Utils.EnableIpv6(self.mock_setup, ['A', 'B'], self.mock_logger)
     expected_calls = [
-        mock.call.call(['A', 'B'], mock.ANY, dhclient_script='test_script'),
+        mock.call.enable_ra(['A', 'B'], mock.ANY),
+        mock.call.dhclient(['A', 'B'], mock.ANY, dhclient_script=None),
+    ]
+    self.assertEqual(mocks.mock_calls, expected_calls)
+
+  @mock.patch('google_compute_engine.distro_lib.helpers.CallDhclientIpv6')
+  def testDisableIpv6(self, mock_call_dhclient):
+    mocks = mock.Mock()
+    mocks.attach_mock(mock_call_dhclient, 'dhclient')
+
+    utils.Utils.DisableIpv6(self.mock_setup, ['A', 'B'], self.mock_logger)
+    expected_calls = [
+        mock.call.dhclient(['A', 'B'], mock.ANY, None, release_lease=True),
     ]
     self.assertEqual(mocks.mock_calls, expected_calls)