From 9642c1171c3034394c57b62b4c8d7e5661122178 Mon Sep 17 00:00:00 2001
From: Ondrej Nosek <onosek@redhat.com>
Date: Mon, 4 Dec 2017 15:15:20 +0100
Subject: [PATCH] add ability to specify ostree ref in OSTREE phase - update

Additionally ostree_ref (if parameter is given) should be placed into treefile.

Relates: https://pagure.io/pungi/issue/777

Signed-off-by: Ondrej Nosek <onosek@redhat.com>
---
 pungi/ostree/tree.py        | 21 ++++++++++++++++-----
 pungi/ostree/utils.py       |  8 +++++++-
 tests/test_ostree_script.py | 31 +++++++++++++++++++++++++++++++
 3 files changed, 54 insertions(+), 6 deletions(-)

diff --git a/pungi/ostree/tree.py b/pungi/ostree/tree.py
index 5e985434..45978a71 100644
--- a/pungi/ostree/tree.py
+++ b/pungi/ostree/tree.py
@@ -85,11 +85,22 @@ class Tree(OSTree):
         self.extra_config = self.args.extra_config
         self.ostree_ref = self.args.ostree_ref
 
-        if self.extra_config:
-            self.extra_config = json.load(open(self.extra_config, 'r'))
-            repos = self.extra_config.get('repo', [])
-            keep_original_sources = self.extra_config.get('keep_original_sources', False)
-            tweak_treeconf(self.treefile, source_repos=repos, keep_original_sources=keep_original_sources)
+        if self.extra_config or self.ostree_ref:
+            if self.extra_config:
+                self.extra_config = json.load(open(self.extra_config, 'r'))
+                repos = self.extra_config.get('repo', [])
+                keep_original_sources = self.extra_config.get('keep_original_sources', False)
+            else:
+                # missing extra_config mustn't affect tweak_treeconf call
+                repos = []
+                keep_original_sources = True
+
+            update_dict = {}
+            if self.ostree_ref:
+                # override ref value in treefile
+                update_dict['ref'] = self.ostree_ref
+
+            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 c832df3b..4c52fd28 100644
--- a/pungi/ostree/utils.py
+++ b/pungi/ostree/utils.py
@@ -79,10 +79,12 @@ def _write_repofile(path, name, repo):
         f.write("gpgcheck=%s\n" % gpgcheck)
 
 
-def tweak_treeconf(treeconf, source_repos=None, keep_original_sources=False):
+def tweak_treeconf(treeconf, source_repos=None, keep_original_sources=False, update_dict=None):
     """
     Update tree config file by adding new repos, and remove existing repos
     from the tree config file if 'keep_original_sources' is not enabled.
+    Additionally, other values can be passed to method by 'update_dict' parameter to
+    update treefile content.
     """
     # add this timestamp to repo name to get unique repo filename and repo name
     # should be safe enough
@@ -108,6 +110,10 @@ def tweak_treeconf(treeconf, source_repos=None, keep_original_sources=False):
     else:
         treeconf_content['repos'] = repos
 
+    # update content with config values from dictionary (for example 'ref')
+    if isinstance(update_dict, dict):
+        treeconf_content.update(update_dict)
+
     # update tree config to add new repos
     with open(treeconf, 'w') as f:
         json.dump(treeconf_content, f, indent=4)
diff --git a/tests/test_ostree_script.py b/tests/test_ostree_script.py
index e1c9ae3d..a2383df4 100644
--- a/tests/test_ostree_script.py
+++ b/tests/test_ostree_script.py
@@ -133,6 +133,37 @@ class OstreeTreeScriptTest(helpers.PungiTestCase):
                         self.topdir + '/fedora-atomic-docker-host.json'],
                        logfile=self.topdir + '/logs/Atomic/create-ostree-repo.log', show_cmd=True, stdout=True)])
 
+    @mock.patch('kobo.shortcuts.run')
+    def test_ostree_ref(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.json')
+
+        with open(treefile, 'r') as f:
+            treefile_content = json.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, 'r') as f:
+            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('pungi.ostree.utils.datetime')
     @mock.patch('kobo.shortcuts.run')
     def test_extra_config_with_extra_repos(self, run, time):