X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/b8261deb6e1035134d7ce8f004cc49deeed1bd27..f486405f195083289f543b5524c4059b01f49026:/services/nodemanager/tests/test_computenode_driver_gce.py

diff --git a/services/nodemanager/tests/test_computenode_driver_gce.py b/services/nodemanager/tests/test_computenode_driver_gce.py
index dd0d5c686a..b9d7ee9fd0 100644
--- a/services/nodemanager/tests/test_computenode_driver_gce.py
+++ b/services/nodemanager/tests/test_computenode_driver_gce.py
@@ -14,6 +14,20 @@ from . import testutil
 class GCEComputeNodeDriverTestCase(testutil.DriverTestMixin, unittest.TestCase):
     TEST_CLASS = gce.ComputeNodeDriver
 
+    def setUp(self):
+        super(GCEComputeNodeDriverTestCase, self).setUp()
+        self.driver_mock().list_images.return_value = [
+            testutil.cloud_object_mock('testimage', selfLink='image-link')]
+        self.driver_mock().ex_list_disktypes.return_value = [
+            testutil.cloud_object_mock(name, selfLink=name + '-link')
+            for name in ['pd-standard', 'pd-ssd', 'local-ssd']]
+        self.driver_mock.reset_mock()
+
+    def new_driver(self, auth_kwargs={}, list_kwargs={}, create_kwargs={}):
+        create_kwargs.setdefault('image', 'testimage')
+        return super(GCEComputeNodeDriverTestCase, self).new_driver(
+            auth_kwargs, list_kwargs, create_kwargs)
+
     def test_driver_instantiation(self):
         kwargs = {'user_id': 'foo'}
         driver = self.new_driver(auth_kwargs=kwargs)
@@ -50,6 +64,26 @@ class GCEComputeNodeDriverTestCase(testutil.DriverTestMixin, unittest.TestCase):
         self.assertEqual(['testA', 'testB'],
                          self.driver_mock().create_node.call_args[1]['ex_tags'])
 
+    def test_create_with_two_disks_attached(self):
+        driver = self.new_driver(create_kwargs={'image': 'testimage'})
+        driver.create_node(testutil.MockSize(1), testutil.arvados_node_mock())
+        create_disks = self.driver_mock().create_node.call_args[1].get(
+            'ex_disks_gce_struct', [])
+        self.assertEqual(2, len(create_disks))
+        self.assertTrue(create_disks[0].get('autoDelete'))
+        self.assertTrue(create_disks[0].get('boot'))
+        self.assertEqual('PERSISTENT', create_disks[0].get('type'))
+        init_params = create_disks[0].get('initializeParams', {})
+        self.assertEqual('pd-standard-link', init_params.get('diskType'))
+        self.assertEqual('image-link', init_params.get('sourceImage'))
+        # Our node images expect the SSD to be named `tmp` to find and mount it.
+        self.assertEqual('tmp', create_disks[1].get('deviceName'))
+        self.assertTrue(create_disks[1].get('autoDelete'))
+        self.assertFalse(create_disks[1].get('boot', 'unset'))
+        self.assertEqual('SCRATCH', create_disks[1].get('type'))
+        init_params = create_disks[1].get('initializeParams', {})
+        self.assertEqual('local-ssd-link', init_params.get('diskType'))
+
     def test_list_nodes_requires_tags_match(self):
         # A node matches if our list tags are a subset of the node's tags.
         # Test behavior with no tags, no match, partial matches, different