From 685ffd49561bb92971f6b76e4690b86d7d6ecc0f Mon Sep 17 00:00:00 2001 From: Scott Moser Date: Fri, 19 Aug 2016 13:58:52 -0400 Subject: 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. --- cloudinit/atomic_helper.py | 20 ++++++++----- cloudinit/cmd/main.py | 12 ++++---- cloudinit/dhclient_hook.py | 4 +-- tests/unittests/helpers.py | 14 +++++++++ tests/unittests/test_atomic_helper.py | 54 +++++++++++++++++++++++++++++++++++ 5 files changed, 89 insertions(+), 15 deletions(-) create mode 100644 tests/unittests/test_atomic_helper.py 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)) -- cgit v1.2.3