freezer: clean freeze YAML profile on restore

doc/topics/releases/neon.rst

@@ -389,6 +389,10 @@ Module Changes
 - The :py:func:`chocolatey.unbootstrap <salt.modules.chocolatey.unbootstrap>` function
   has been added to uninstall Chocolatey.
 
+- The :py:func:`freezer.restore <salt.modules.freezer.restore>` function provides a
+  new parameter, `clean`, that will remove the frozen data after a successful recovery.
+
+
 Runner Changes
 ==============
 

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

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)