Support distro configs for instance defaults. (#286)

2 files changed, 49 insertions, 2 deletions

diff --git a/google_compute_engine/instance_setup/instance_config.py b/google_compute_engine/instance_setup/instance_config.py
index 2d39283..06c1cb8 100644
--- a/google_compute_engine/instance_setup/instance_config.py
+++ b/google_compute_engine/instance_setup/instance_config.py
@@ -31,7 +31,8 @@ class InstanceConfig(config_manager.ConfigManager):
   """Creates a defaults config file for instance configuration."""
 
   instance_config = '/etc/default/instance_configs.cfg'
-  instance_config_template = '/etc/default/instance_configs.cfg.template'
+  instance_config_distro = '%s.distro' % instance_config
+  instance_config_template = '%s.template' % instance_config
   instance_config_script = os.path.abspath(__file__)
   instance_config_header = (
       'This file is automatically created at boot time by the %s script. Do '
@@ -75,13 +76,24 @@ class InstanceConfig(config_manager.ConfigManager):
     """
     self.instance_config_header %= (
         self.instance_config_script, self.instance_config_template)
+    # User provided instance configs should always take precedence.
     super(InstanceConfig, self).__init__(
         config_file=self.instance_config_template,
         config_header=self.instance_config_header)
 
+    # Use the settings in an instance config file if one exists. If a config
+    # file does not already exist, try to use the distro provided defaults. If
+    # no file exists, use the default configuration settings.
     if os.path.exists(self.instance_config):
+      instance_config = self.instance_config
+    elif os.path.exists(self.instance_config_distro):
+      instance_config = self.instance_config_distro
+    else:
+      instance_config = None
+
+    if instance_config:
       config = parser.SafeConfigParser()
-      config.read(self.instance_config)
+      config.read(instance_config)
       defaults = dict((s, dict(config.items(s))) for s in config.sections())
     else:
       defaults = self.instance_config_options
diff --git a/google_compute_engine/instance_setup/tests/instance_config_test.py b/google_compute_engine/instance_setup/tests/instance_config_test.py
index 6e98f19..342d733 100644
--- a/google_compute_engine/instance_setup/tests/instance_config_test.py
+++ b/google_compute_engine/instance_setup/tests/instance_config_test.py
@@ -24,6 +24,7 @@ class InstanceConfigTest(unittest.TestCase):
 
   def setUp(self):
     instance_config.InstanceConfig.instance_config = 'config'
+    instance_config.InstanceConfig.instance_config_distro = 'distro'
     instance_config.InstanceConfig.instance_config_template = 'template'
     instance_config.InstanceConfig.instance_config_script = '/tmp/test.py'
     instance_config.InstanceConfig.instance_config_header = '%s %s'
@@ -56,6 +57,7 @@ class InstanceConfigTest(unittest.TestCase):
         mock.call.init(
             config_file='template', config_header='/tmp/test.py template'),
         mock.call.exists('config'),
+        mock.call.exists('distro'),
         mock.call.set('first', 'a', 'false', overwrite=False),
         mock.call.set('second', 'b', 'true', overwrite=False),
         mock.call.set('third', 'c', '1', overwrite=False),
@@ -95,6 +97,39 @@ class InstanceConfigTest(unittest.TestCase):
     ]
     self.assertEqual(mocks.mock_calls, expected_calls)
 
+  @mock.patch('google_compute_engine.instance_setup.instance_config.os.path.exists')
+  @mock.patch('google_compute_engine.instance_setup.instance_config.parser')
+  @mock.patch('google_compute_engine.instance_setup.instance_config.config_manager.ConfigManager.SetOption')
+  @mock.patch('google_compute_engine.instance_setup.instance_config.config_manager.ConfigManager.__init__')
+  def testInstanceConfigDistroExists(self, mock_init, mock_set, mock_parser,
+                                     mock_exists):
+    mock_config = mock.create_autospec(instance_config.parser.SafeConfigParser)
+    mock_config.read = mock.Mock()
+    mock_config.sections = mock.Mock()
+    mock_config.sections.return_value = ['a', 'b']
+    mock_config.items = lambda key: {'key: %s' % key: 'value: %s' % key}
+    mock_parser.SafeConfigParser.return_value = mock_config
+    mocks = mock.Mock()
+    mocks.attach_mock(mock_init, 'init')
+    mocks.attach_mock(mock_set, 'set')
+    mocks.attach_mock(mock_parser, 'parser')
+    mocks.attach_mock(mock_exists, 'exists')
+    mock_exists.side_effect = [False, True]
+
+    instance_config.InstanceConfig()
+    expected_calls = [
+        mock.call.init(
+            config_file='template', config_header='/tmp/test.py template'),
+        mock.call.exists('config'),
+        mock.call.exists('distro'),
+        mock.call.parser.SafeConfigParser(),
+        mock.call.parser.SafeConfigParser().read('distro'),
+        mock.call.parser.SafeConfigParser().sections(),
+        mock.call.set('a', 'key: a', 'value: a', overwrite=False),
+        mock.call.set('b', 'key: b', 'value: b', overwrite=False),
+    ]
+    self.assertEqual(mocks.mock_calls, expected_calls)
+
   @mock.patch('google_compute_engine.instance_setup.instance_config.config_manager.ConfigManager.WriteConfig')
   def testWriteConfig(self, mock_write):
     mock_config = instance_config.InstanceConfig()