diff --git a/doc/contributing.rst b/doc/contributing.rst
index 4aa072fe..7b4dc790 100644
--- a/doc/contributing.rst
+++ b/doc/contributing.rst
@@ -32,6 +32,7 @@ These packages will have to installed:
  * python-lxml
  * python2-multilib
  * python-productmd
+ * PyYAML
  * repoview
  * rpm-devel
  * syslinux
@@ -59,7 +60,7 @@ packages above as they are used by calling an executable. ::
     $ for pkg in _deltarpm krbV _selinux deltarpm sqlitecachec _sqlitecache; do ln -vs "$(deactivate && python -c 'import os, '$pkg'; print('$pkg'.__file__)')" "$(virtualenvwrapper_get_site_packages_dir)"; done
     $ pip install -U pip
     $ PYCURL_SSL_LIBRARY=nss pip install pycurl --no-binary :all:
-    $ pip install beanbag jsonschema 'kobo>=0.6.0' lockfile lxml mock nose nose-cov productmd pyopenssl python-multilib requests requests-kerberos setuptools sphinx ordered_set koji
+    $ pip install beanbag jsonschema 'kobo>=0.6.0' lockfile lxml mock nose nose-cov productmd pyopenssl python-multilib requests requests-kerberos setuptools sphinx ordered_set koji PyYAML
 
 Now you should be able to run all existing tests.
 
diff --git a/pungi/ostree/tree.py b/pungi/ostree/tree.py
index d430a467..fceca6f0 100644
--- a/pungi/ostree/tree.py
+++ b/pungi/ostree/tree.py
@@ -101,7 +101,12 @@ class Tree(OSTree):
                 # override ref value in treefile
                 update_dict['ref'] = self.ostree_ref
 
-            self.treefile = tweak_treeconf(self.treefile, source_repos=repos, keep_original_sources=keep_original_sources, update_dict=update_dict)
+            self.treefile = tweak_treeconf(
+                self.treefile,
+                source_repos=repos,
+                keep_original_sources=keep_original_sources,
+                update_dict=update_dict
+            )
 
         self.commitid_file = make_log_file(self.logdir, 'commitid')
 
diff --git a/pungi/ostree/utils.py b/pungi/ostree/utils.py
index 39d3a379..d76422b0 100644
--- a/pungi/ostree/utils.py
+++ b/pungi/ostree/utils.py
@@ -14,10 +14,12 @@
 # along with this program; if not, see <https://gnu.org/licenses/>.
 
 
-import yaml
 import json
 import logging
 import os
+import shutil
+
+import yaml
 
 from pungi.arch_utils import getBaseArch
 from pungi.util import makedirs
@@ -75,12 +77,12 @@ def tweak_treeconf(treeconf, source_repos=None, keep_original_sources=False, upd
     """
 
     # backup the old tree config
-    os.rename(treeconf, '{0}.bak'.format(treeconf))
+    shutil.copy2(treeconf, '{0}.bak'.format(treeconf))
 
     treeconf_dir = os.path.dirname(treeconf)
     with open(treeconf, 'r') as f:
-        # rpm-ostree now supports YAML: https://github.com/projectatomic/rpm-ostree/pull/1377
-        # but we'll end up converting it to JSON.
+        # rpm-ostree now supports YAML, but we'll end up converting it to JSON.
+        # https://github.com/projectatomic/rpm-ostree/pull/1377
         if treeconf.endswith('.yaml'):
             treeconf_content = yaml.load(f)
             treeconf = treeconf.replace('.yaml', '.json')
diff --git a/tests/test_ostree_script.py b/tests/test_ostree_script.py
index ad98052d..ce211c3e 100644
--- a/tests/test_ostree_script.py
+++ b/tests/test_ostree_script.py
@@ -2,12 +2,13 @@
 # -*- coding: utf-8 -*-
 
 
-import mock
-
-import os
 import json
+import os
 import sys
 
+import mock
+import yaml
+
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'bin'))
 
@@ -21,6 +22,9 @@ class OstreeTreeScriptTest(helpers.PungiTestCase):
         helpers.touch(os.path.join(path, 'fedora-atomic-docker-host.json'),
                       json.dumps({'ref': 'fedora-atomic/25/x86_64',
                                   'repos': ['fedora-rawhide', 'fedora-24', 'fedora-23']}))
+        helpers.touch(os.path.join(path, 'fedora-atomic-docker-host.yaml'),
+                      yaml.dump({'ref': 'fedora-atomic/25/x86_64',
+                                 'repos': ['fedora-rawhide', 'fedora-24', 'fedora-23']}))
         helpers.touch(os.path.join(path, 'fedora-rawhide.repo'),
                       '[fedora-rawhide]\nmirrorlist=mirror-mirror-on-the-wall')
         helpers.touch(os.path.join(path, 'fedora-24.repo'),
@@ -163,6 +167,39 @@ class OstreeTreeScriptTest(helpers.PungiTestCase):
         # repos should stay unchanged
         self.assertEqual(original_repos, new_repos)
 
+    @mock.patch('kobo.shortcuts.run')
+    def test_run_with_yaml_file(self, run):
+        repo = os.path.join(self.topdir, 'atomic')
+
+        self._make_dummy_config_dir(self.topdir)
+        treefile = os.path.join(self.topdir, 'fedora-atomic-docker-host.yaml')
+
+        with open(treefile, 'r') as f:
+            # Read initial content from YAML file
+            treefile_content = yaml.load(f)
+        original_repos = treefile_content['repos']
+        original_ref = treefile_content['ref']
+        replacing_ref = original_ref + '-changed'
+
+        ostree.main([
+            'tree',
+            '--repo=%s' % repo,
+            '--log-dir=%s' % os.path.join(self.topdir, 'logs', 'Atomic'),
+            '--treefile=%s' % treefile,
+            '--ostree-ref=%s' % replacing_ref,
+        ])
+
+        with open(treefile.replace(".yaml", ".json"), 'r') as f:
+            # There is now a tweaked JSON file
+            treefile_content = json.load(f)
+        new_repos = treefile_content['repos']
+        new_ref = treefile_content['ref']
+
+        # ref value in treefile should be overrided with new ref
+        self.assertEqual(replacing_ref, new_ref)
+        # repos should stay unchanged
+        self.assertEqual(original_repos, new_repos)
+
     @mock.patch('kobo.shortcuts.run')
     def test_force_new_commit(self, run):
         repo = os.path.join(self.topdir, 'atomic')