Minor cleanups to atomic_helper and add unit tests.

Change atomic_helper.write_file to have same same signature as write_file.
Add some simple unit tests for atomic_helper.

5 files changed, 89 insertions, 15 deletions

diff --git a/cloudinit/atomic_helper.py b/cloudinit/atomic_helper.py
index 15319f71..a3cfd942 100644
--- a/cloudinit/atomic_helper.py
+++ b/cloudinit/atomic_helper.py
@@ -5,21 +5,27 @@ import json
 import os
 import tempfile
 
+_DEF_PERMS = 0o644
 
-def atomic_write_file(path, content, mode='w'):
+
+def write_file(filename, content, mode=_DEF_PERMS, omode="wb"):
+    # open filename in mode 'omode', write content, set permissions to 'mode'
     tf = None
     try:
-        tf = tempfile.NamedTemporaryFile(dir=os.path.dirname(path),
-                                         delete=False, mode=mode)
+        tf = tempfile.NamedTemporaryFile(dir=os.path.dirname(filename),
+                                         delete=False, mode=omode)
         tf.write(content)
         tf.close()
-        os.rename(tf.name, path)
+        os.chmod(tf.name, mode)
+        os.rename(tf.name, filename)
     except Exception as e:
         if tf is not None:
             os.unlink(tf.name)
         raise e
 
 
-def atomic_write_json(path, data):
-    return atomic_write_file(path, json.dumps(data, indent=1,
-                                              sort_keys=True) + "\n")
+def write_json(filename, data, mode=_DEF_PERMS):
+    # dump json representation of data to file filename.
+    return write_file(
+        filename, json.dumps(data, indent=1, sort_keys=True) + "\n",
+        omode="w", mode=mode)
diff --git a/cloudinit/cmd/main.py b/cloudinit/cmd/main.py
index ba22b168..83eb02c9 100644
--- a/cloudinit/cmd/main.py
+++ b/cloudinit/cmd/main.py
@@ -46,7 +46,7 @@ from cloudinit.reporting import events
 from cloudinit.settings import (PER_INSTANCE, PER_ALWAYS, PER_ONCE,
                                 CLOUD_CONFIG)
 
-from cloudinit.atomic_helper import atomic_write_json
+from cloudinit import atomic_helper
 
 from cloudinit.dhclient_hook import LogDhclient
 
@@ -513,7 +513,7 @@ def status_wrapper(name, args, data_d=None, link_d=None):
     v1['stage'] = mode
     v1[mode]['start'] = time.time()
 
-    atomic_write_json(status_path, status)
+    atomic_helper.write_json(status_path, status)
     util.sym_link(os.path.relpath(status_path, link_d), status_link,
                   force=True)
 
@@ -536,7 +536,7 @@ def status_wrapper(name, args, data_d=None, link_d=None):
     v1[mode]['finished'] = time.time()
     v1['stage'] = None
 
-    atomic_write_json(status_path, status)
+    atomic_helper.write_json(status_path, status)
 
     if mode == "modules-final":
         # write the 'finished' file
@@ -545,9 +545,9 @@ def status_wrapper(name, args, data_d=None, link_d=None):
             if v1[m]['errors']:
                 errors.extend(v1[m].get('errors', []))
 
-        atomic_write_json(result_path,
-                          {'v1': {'datasource': v1['datasource'],
-                                  'errors': errors}})
+        atomic_helper.write_json(
+            result_path, {'v1': {'datasource': v1['datasource'],
+                          'errors': errors}})
         util.sym_link(os.path.relpath(result_path, link_d), result_link,
                       force=True)
 
diff --git a/cloudinit/dhclient_hook.py b/cloudinit/dhclient_hook.py
index 9dcbe39c..82cb1855 100644
--- a/cloudinit/dhclient_hook.py
+++ b/cloudinit/dhclient_hook.py
@@ -3,7 +3,7 @@
 
 import os
 
-from cloudinit.atomic_helper import atomic_write_json
+from cloudinit import atomic_helper
 from cloudinit import log as logging
 from cloudinit import stages
 
@@ -46,5 +46,5 @@ class LogDhclient(object):
         envs = os.environ
         if self.hook_file is None:
             return
-        atomic_write_json(self.hook_file, self.get_vals(envs))
+        atomic_helper.write_json(self.hook_file, self.get_vals(envs))
         LOG.debug("Wrote dhclient options in %s", self.hook_file)
diff --git a/tests/unittests/helpers.py b/tests/unittests/helpers.py
index de2cf638..1cdc05a1 100644
--- a/tests/unittests/helpers.py
+++ b/tests/unittests/helpers.py
@@ -252,6 +252,20 @@ class HttprettyTestCase(TestCase):
         super(HttprettyTestCase, self).tearDown()
 
 
+class TempDirTestCase(TestCase):
+    # provide a tempdir per class, not per test.
+    def setUp(self):
+        super(TempDirTestCase, self).setUp()
+        self.tmp = tempfile.mkdtemp()
+        self.addCleanup(shutil.rmtree, self.tmp)
+
+    def tmp_path(self, path):
+        if path.startswith(os.path.sep):
+            path = "." + path
+
+        return os.path.normpath(os.path.join(self.tmp, path))
+
+
 def populate_dir(path, files):
     if not os.path.exists(path):
         os.makedirs(path)
diff --git a/tests/unittests/test_atomic_helper.py b/tests/unittests/test_atomic_helper.py
new file mode 100644
index 00000000..feb81551
--- /dev/null
+++ b/tests/unittests/test_atomic_helper.py
@@ -0,0 +1,54 @@
+import json
+import os
+import stat
+
+from cloudinit import atomic_helper
+
+from . import helpers
+
+
+class TestAtomicHelper(helpers.TempDirTestCase):
+    def test_basic_usage(self):
+        """write_file takes bytes if no omode."""
+        path = self.tmp_path("test_basic_usage")
+        contents = b"Hey there\n"
+        atomic_helper.write_file(path, contents)
+        self.check_file(path, contents)
+
+    def test_string(self):
+        """write_file can take a string with mode w."""
+        path = self.tmp_path("test_string")
+        contents = "Hey there\n"
+        atomic_helper.write_file(path, contents, omode="w")
+        self.check_file(path, contents, omode="r")
+
+    def test_file_permissions(self):
+        """write_file with mode 400 works correctly."""
+        path = self.tmp_path("test_file_permissions")
+        contents = b"test_file_perms"
+        atomic_helper.write_file(path, contents, mode=0o400)
+        self.check_file(path, contents, perms=0o400)
+
+    def test_write_json(self):
+        """write_json output is readable json."""
+        path = self.tmp_path("test_write_json")
+        data = {'key1': 'value1', 'key2': ['i1', 'i2']}
+        atomic_helper.write_json(path, data)
+        with open(path, "r") as fp:
+            found = json.load(fp)
+        self.assertEqual(data, found)
+        self.check_perms(path, 0o644)
+
+    def check_file(self, path, content, omode=None, perms=0o644):
+        if omode is None:
+            omode = "rb"
+        self.assertTrue(os.path.exists(path))
+        self.assertTrue(os.path.isfile(path))
+        with open(path, omode) as fp:
+            found = fp.read()
+            self.assertEqual(content, found)
+        self.check_perms(path, perms)
+
+    def check_perms(self, path, perms):
+        file_stat = os.stat(path)
+        self.assertEqual(perms, stat.S_IMODE(file_stat.st_mode))