diff --git a/doc/topics/releases/neon.rst b/doc/topics/releases/neon.rst index 9ca10428e27b..b44030a7cb07 100644 --- a/doc/topics/releases/neon.rst +++ b/doc/topics/releases/neon.rst @@ -389,6 +389,10 @@ Module Changes - The :py:func:`chocolatey.unbootstrap ` function has been added to uninstall Chocolatey. +- The :py:func:`freezer.restore ` function provides a + new parameter, `clean`, that will remove the frozen data after a successful recovery. + + Runner Changes ============== diff --git a/salt/modules/freezer.py b/salt/modules/freezer.py index 6554e03815c5..4d6c907506c3 100644 --- a/salt/modules/freezer.py +++ b/salt/modules/freezer.py @@ -164,13 +164,13 @@ def freeze(name=None, force=False, **kwargs): safe_kwargs = clean_kwargs(**kwargs) pkgs = __salt__['pkg.list_pkgs'](**safe_kwargs) repos = __salt__['pkg.list_repos'](**safe_kwargs) - for name, content in zip(_paths(name), (pkgs, repos)): - with fopen(name, 'w') as fp: + for fname, content in zip(_paths(name), (pkgs, repos)): + with fopen(fname, 'w') as fp: json.dump(content, fp) return True -def restore(name=None, **kwargs): +def restore(name=None, clean=False, **kwargs): ''' Make sure that the system contains the packages and repos from a frozen state. @@ -191,6 +191,11 @@ def restore(name=None, **kwargs): name Name of the frozen state. Optional. + clean + If True remove the frozen information YAML from the cache + + .. version-added:: Neon + CLI Example: .. code-block:: bash @@ -204,8 +209,8 @@ def restore(name=None, **kwargs): frozen_pkgs = {} frozen_repos = {} - for name, content in zip(_paths(name), (frozen_pkgs, frozen_repos)): - with fopen(name) as fp: + for fname, content in zip(_paths(name), (frozen_pkgs, frozen_repos)): + with fopen(fname) as fp: content.update(json.load(fp)) # The ordering of removing or adding packages and repos can be @@ -292,4 +297,9 @@ def restore(name=None, **kwargs): log.error(msg, repo, e) res['comment'].append(msg % (repo, e)) + # Clean the cached YAML files + if clean and not res['comment']: + for fname in _paths(name): + os.remove(fname) + return res diff --git a/tests/unit/modules/test_freezer.py b/tests/unit/modules/test_freezer.py index b232e93f29c4..70d315c17a38 100644 --- a/tests/unit/modules/test_freezer.py +++ b/tests/unit/modules/test_freezer.py @@ -296,3 +296,29 @@ def test_restore_remove_extra_repo(self, status, fopen, load): salt_mock['pkg.del_repo'].assert_called_once() fopen.assert_called() load.assert_called() + + @patch('os.remove') + @patch('salt.utils.json.load') + @patch('salt.modules.freezer.fopen') + @patch('salt.modules.freezer.status') + def test_restore_clean_yml(self, status, fopen, load, remove): + ''' + Test to restore an old state + ''' + status.return_value = True + salt_mock = { + 'pkg.list_pkgs': MagicMock(return_value={}), + 'pkg.list_repos': MagicMock(return_value={}), + 'pkg.install': MagicMock(), + } + with patch.dict(freezer.__salt__, salt_mock): + self.assertEqual(freezer.restore(clean=True), { + 'pkgs': {'add': [], 'remove': []}, + 'repos': {'add': [], 'remove': []}, + 'comment': [], + }) + salt_mock['pkg.list_pkgs'].assert_called() + salt_mock['pkg.list_repos'].assert_called() + fopen.assert_called() + load.assert_called() + self.assertEqual(remove.call_count, 2)