1 files changed, 40 insertions, 2 deletions
diff --git a/packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py b/packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py
index a418219..5c6e513 100644
--- a/packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py
+++ b/packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py
@@ -46,7 +46,10 @@ class InstanceSetupTest(unittest.TestCase):
     mock_logger_instance = mock.Mock()
     mock_logger.Logger.return_value = mock_logger_instance
     mock_watcher_instance = mock.Mock()
-    mock_watcher_instance.GetMetadata.return_value = {'hello': 'world'}
+    mock_watcher_instance.GetMetadata.return_value = {
+        'hello': 'world',
+        'instance': {'machineType': 'fake'},
+    }
     mock_watcher.MetadataWatcher.return_value = mock_watcher_instance
     mock_config_instance = mock.Mock()
     mock_config_instance.GetOptionBool.return_value = True
@@ -79,6 +82,7 @@ class InstanceSetupTest(unittest.TestCase):
         mock.call.config.InstanceConfig().GetOptionBool(
             'InstanceSetup', 'set_boto_config'),
         mock.call.setup._SetupBotoConfig(),
+        mock.call.setup._DisableOvercommit(),
         # Setup for local SSD.
         mock.call.config.InstanceConfig().GetOptionBool(
             'InstanceSetup', 'optimize_local_ssd'),
@@ -91,7 +95,10 @@ class InstanceSetupTest(unittest.TestCase):
         mock.call.config.InstanceConfig().WriteConfig(),
     ]
     self.assertEqual(mocks.mock_calls, expected_calls)
-    self.assertEqual(mock_setup.metadata_dict, {'hello': 'world'})
+    self.assertEqual(mock_setup.metadata_dict, {
+        'hello': 'world',
+        'instance':{'machineType': 'fake'},
+    })
 
   @mock.patch('google_compute_engine.instance_setup.instance_setup.instance_config')
   @mock.patch('google_compute_engine.instance_setup.instance_setup.metadata_watcher')
@@ -421,6 +428,37 @@ class InstanceSetupTest(unittest.TestCase):
     instance_setup.InstanceSetup._SetupBotoConfig(self.mock_setup)
     self.mock_logger.warning.assert_called_once_with('Test Error')
 
+  @mock.patch('google_compute_engine.instance_setup.instance_setup.subprocess')
+  def testDisableOvercommitNonE2(self, mock_subprocess):
+    self.mock_setup.metadata_dict = {
+        'instance': {
+            'machineType': 'projects/00000000000/machineTypes/n1-standard-1',
+        }
+    }
+    instance_setup.InstanceSetup._DisableOvercommit(self.mock_setup)
+    mock_subprocess.call.assert_not_called()
+
+  @mock.patch('google_compute_engine.instance_setup.instance_setup.subprocess')
+  def testDisableOvercommitE2(self, mock_subprocess):
+    self.mock_setup.metadata_dict = {
+        'instance': {
+            'machineType': 'projects/00000000000/machineTypes/e2-standard-1',
+        }
+    }
+    instance_setup.InstanceSetup._DisableOvercommit(self.mock_setup)
+    mock_subprocess.call.assert_called_once_with(
+        ['sysctl', 'vm.overcommit_memory=1'])
+
+  @mock.patch('google_compute_engine.instance_setup.instance_setup.subprocess')
+  def testDisableOvercommitBSD(self, mock_subprocess):
+    self.mock_setup.metadata_dict = {
+        'instance': {
+            'machineType': 'projects/00000000000/machineTypes/e2-standard-1',
+        }
+    }
+    instance_setup.InstanceSetup._DisableOvercommit(self.mock_setup, 'bsd')
+    mock_subprocess.call.assert_not_called()
+
 
 if __name__ == '__main__':
   unittest.main()