From a015473d4e56af9aae01ef465887ec6ad9e4bf0c Mon Sep 17 00:00:00 2001 From: Steve Hajducko Date: Wed, 2 Dec 2015 11:36:46 -0800 Subject: [PATCH] Add etcd update function Adds a update function that takes a dict and sets multiple etcd keys. --- salt/modules/etcd_mod.py | 53 ++++++++++++++++++ salt/utils/etcd_util.py | 31 +++++++++++ tests/unit/modules/etcd_mod_test.py | 30 ++++++++++ tests/unit/utils/etcd_util_test.py | 85 +++++++++++++++++++++++++++++ 4 files changed, 199 insertions(+) diff --git a/salt/modules/etcd_mod.py b/salt/modules/etcd_mod.py index b4d6e36152c4..f3c9ca4c445e 100644 --- a/salt/modules/etcd_mod.py +++ b/salt/modules/etcd_mod.py @@ -101,6 +101,59 @@ def set_(key, value, profile=None, ttl=None, directory=False): return client.set(key, value, ttl=ttl, directory=directory) +def update(fields, path='', profile=None): + ''' + .. versionadded:: Boron + + Sets a dictionary of values in one call. Useful for large updates + in syndic environments. The dictionary can contain a mix of formats + such as: + + .. code-block:: python + + { + '/some/example/key': 'bar', + '/another/example/key': 'baz' + } + + Or it may be a straight dictionary, which will be flattened to look + like the above format: + + .. code-block:: python + + { + 'some': { + 'example': { + 'key': 'bar' + } + }, + 'another': { + 'example': { + 'key': 'baz' + } + } + } + + You can even mix the two formats and it will be flattened to the first + format. Leading and trailing '/' will be removed. + + Empty directories can be created by setting the value of the key to an + empty dictionary. + + The 'path' parameter will optionally set the root of the path to use. + + CLI Example: + + .. code-block:: bash + + salt myminion etcd.update "{'/path/to/key': 'baz', '/another/key': 'bar'}" + salt myminion etcd.update "{'/path/to/key': 'baz', '/another/key': 'bar'}" profile=my_etcd_config + salt myminion etcd.update "{'/path/to/key': 'baz', '/another/key': 'bar'}" path='/some/root' + ''' + client = __utils__['etcd_util.get_conn'](__opts__, profile) + return client.update(fields, path) + + def watch(key, recurse=False, profile=None, timeout=0, index=None): ''' .. versionadded:: Boron diff --git a/salt/utils/etcd_util.py b/salt/utils/etcd_util.py index 7ed6eb00874e..6eeaba1de3f5 100644 --- a/salt/utils/etcd_util.py +++ b/salt/utils/etcd_util.py @@ -187,6 +187,37 @@ def read(self, key, recursive=False, wait=False, timeout=None, waitIndex=None): raise return result + def _flatten(self, data, path=''): + if len(data.keys()) == 0: + return {path: {}} + path = path.strip('/') + flat = {} + for k, v in data.iteritems(): + k = k.strip('/') + if path: + p = '/{0}/{1}'.format(path, k) + else: + p = '/{0}'.format(k) + if isinstance(v, dict): + ret = self._flatten(v, p) + flat.update(ret) + else: + flat[p] = v + return flat + + def update(self, fields, path=''): + if not isinstance(fields, dict): + log.error('etcd.update: fields is not type dict') + return None + fields = self._flatten(fields, path) + keys = {} + for k, v in fields.iteritems(): + is_dir = False + if isinstance(v, dict): + is_dir = True + keys[k] = self.write(k, v, directory=is_dir) + return keys + def set(self, key, value, ttl=None, directory=False): return self.write(key, value, ttl=ttl, directory=directory) diff --git a/tests/unit/modules/etcd_mod_test.py b/tests/unit/modules/etcd_mod_test.py index 2bd4f498f4b2..f3f59be93669 100644 --- a/tests/unit/modules/etcd_mod_test.py +++ b/tests/unit/modules/etcd_mod_test.py @@ -85,6 +85,36 @@ def test_set(self): self.instance.set.side_effect = Exception self.assertRaises(Exception, etcd_mod.set_, 'err', 'stack') + # 'update' function tests: 1 + + def test_update(self): + ''' + Test if can set multiple keys in etcd + ''' + with patch.dict(etcd_mod.__utils__, {'etcd_util.get_conn': self.EtcdClientMock}): + args = { + 'x': { + 'y': { + 'a': '1', + 'b': '2', + } + }, + 'z': '4', + 'd': {}, + } + + result = { + '/some/path/x/y/a': '1', + '/some/path/x/y/b': '2', + '/some/path/z': '4', + '/some/path/d': {}, + } + self.instance.update.return_value = result + self.assertDictEqual(etcd_mod.update(args, path='/some/path'), result) + self.instance.update.assert_called_with(args, '/some/path') + self.assertDictEqual(etcd_mod.update(args), result) + self.instance.update.assert_called_with(args, '') + # 'ls_' function tests: 1 def test_ls(self): diff --git a/tests/unit/utils/etcd_util_test.py b/tests/unit/utils/etcd_util_test.py index 0e1bd60a31e1..917e4d0d9475 100644 --- a/tests/unit/utils/etcd_util_test.py +++ b/tests/unit/utils/etcd_util_test.py @@ -188,6 +188,91 @@ def test_write(self, mock): etcd_client.write.side_effect = Exception self.assertRaises(Exception, client.set, 'some-key', 'some-val') + @patch('etcd.Client', autospec=True) + def test_flatten(self, mock): + client = etcd_util.EtcdClient({}) + some_data = { + '/x/y/a': '1', + 'x': { + 'y': { + 'b': '2' + } + }, + 'm/j/': '3', + 'z': '4', + 'd': {}, + } + + result_path = { + '/test/x/y/a': '1', + '/test/x/y/b': '2', + '/test/m/j': '3', + '/test/z': '4', + '/test/d': {}, + } + + result_nopath = { + '/x/y/a': '1', + '/x/y/b': '2', + '/m/j': '3', + '/z': '4', + '/d': {}, + } + + result_root = { + '/x/y/a': '1', + '/x/y/b': '2', + '/m/j': '3', + '/z': '4', + '/d': {}, + } + + self.assertEqual(client._flatten(some_data, path='/test'), result_path) + self.assertEqual(client._flatten(some_data, path='/'), result_root) + self.assertEqual(client._flatten(some_data), result_nopath) + + @patch('etcd.Client', autospec=True) + def test_update(self, mock): + client = etcd_util.EtcdClient({}) + some_data = { + '/x/y/a': '1', + 'x': { + 'y': { + 'b': '3' + } + }, + 'm/j/': '3', + 'z': '4', + 'd': {}, + } + + result = { + '/test/x/y/a': '1', + '/test/x/y/b': '2', + '/test/m/j': '3', + '/test/z': '4', + '/test/d': True, + } + + flatten_result = { + '/test/x/y/a': '1', + '/test/x/y/b': '2', + '/test/m/j': '3', + '/test/z': '4', + '/test/d': {} + } + client._flatten = MagicMock(return_value=flatten_result) + + self.assertEqual(client.update('/some/key', path='/blah'), None) + + with patch.object(client, 'write', autospec=True) as write_mock: + def write_return(key, val, ttl=None, directory=None): + return result.get(key, None) + write_mock.side_effect = write_return + self.assertDictEqual(client.update(some_data, path='/test'), result) + client._flatten.assert_called_with(some_data, '/test') + self.assertEqual(write_mock.call_count, 5) + @patch('etcd.Client', autospec=True) def test_rm(self, mock): etcd_client = mock.return_value