From f809cac0b2a5ba5a432d48b544966ea1087b503d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lubom=C3=ADr=20Sedl=C3=A1=C5=99?= <lsedlar@redhat.com>
Date: Mon, 19 Nov 2018 11:57:08 +0100
Subject: [PATCH] Cleanup parsing treefile
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

...and also add tests for it.

Signed-off-by: Lubomír Sedlář <lsedlar@redhat.com>
---
 pungi/ostree/utils.py      | 11 +++-----
 tests/test_ostree_utils.py | 52 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 8 deletions(-)
 create mode 100644 tests/test_ostree_utils.py

diff --git a/pungi/ostree/utils.py b/pungi/ostree/utils.py
index 0b74cb23..58368d4a 100644
--- a/pungi/ostree/utils.py
+++ b/pungi/ostree/utils.py
@@ -39,7 +39,6 @@ def get_ref_from_treefile(treefile, arch=None, logger=None):
     the basearch of the architecture we are running on or of the passed in arch.
     """
     logger = logger or logging.getLogger(__name__)
-    ref = None
     if os.path.isfile(treefile):
         with open(treefile, 'r') as f:
             try:
@@ -48,17 +47,13 @@ def get_ref_from_treefile(treefile, arch=None, logger=None):
                 if treefile.endswith(".yaml"):
                     parsed = yaml.safe_load(f)
                 else:
-                    parsed = json.loads(f.read())
-                if arch is None:
-                    basearch = getBaseArch()
-                else:
-                    basearch = getBaseArch(arch)
-                ref = parsed['ref'].replace('${basearch}', basearch)
+                    parsed = json.load(f)
+                return parsed["ref"].replace("${basearch}", getBaseArch(arch))
             except Exception as e:
                 logger.error('Unable to get ref from treefile: %s' % e)
     else:
         logger.error('Unable to open treefile')
-    return ref
+    return None
 
 
 def get_commitid_from_commitid_file(commitid_file):
diff --git a/tests/test_ostree_utils.py b/tests/test_ostree_utils.py
new file mode 100644
index 00000000..bfc699e9
--- /dev/null
+++ b/tests/test_ostree_utils.py
@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+
+import mock
+
+import os
+import sys
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
+
+from tests import helpers
+from pungi.ostree import utils
+
+
+class GetRefFromTreefileTest(helpers.PungiTestCase):
+    def setUp(self):
+        super(GetRefFromTreefileTest, self).setUp()
+        self.json = os.path.join(self.topdir, "treefile.json")
+        self.yaml = os.path.join(self.topdir, "treefile.yaml")
+
+    def test_reads_json(self):
+        helpers.touch(self.json, """{"ref": "master"}""")
+        self.assertEqual(utils.get_ref_from_treefile(self.json), "master")
+
+    def test_reads_yaml(self):
+        helpers.touch(self.yaml, """ref: master""")
+        self.assertEqual(utils.get_ref_from_treefile(self.yaml), "master")
+
+    @mock.patch("pungi.ostree.utils.getBaseArch")
+    def test_replaces_basearch(self, getBaseArch):
+        getBaseArch.return_value = "x86_64"
+        helpers.touch(self.json, """{"ref": "${basearch}/master"}""")
+        self.assertEqual(utils.get_ref_from_treefile(self.json), "x86_64/master")
+
+    @mock.patch("pungi.ostree.utils.getBaseArch")
+    def test_replaces_basearch_for_given_arch(self, getBaseArch):
+        getBaseArch.return_value = "x86_64"
+        helpers.touch(self.json, """{"ref": "${basearch}/master"}""")
+        self.assertEqual(
+            utils.get_ref_from_treefile(self.json, arch="foo"), "x86_64/master"
+        )
+        self.assertEqual(getBaseArch.call_args_list, [mock.call("foo")])
+
+    def test_handles_invalid_json(self):
+        helpers.touch(self.json, """{"ref" "master"}""")
+        self.assertIsNone(utils.get_ref_from_treefile(self.json))
+
+    def test_handles_invalid_yaml(self):
+        helpers.touch(self.yaml, """{ ref\n - master""")
+        self.assertIsNone(utils.get_ref_from_treefile(self.yaml))
+
+    def test_handles_missing_file(self):
+        self.assertIsNone(utils.get_ref_from_treefile(self.json))