From 949e5f7fb5a5344f13ee5e22e1871e42616825c8 Mon Sep 17 00:00:00 2001
From: liaou3 <vincent.liao@canonical.com>
Date: Wed, 3 Jan 2024 16:35:44 +0800
Subject: [PATCH] Add test_image_checker in to checkbox_support

---
 .../checkbox_support/image_checker.py         | 55 ++++++++++++++++
 .../tests/test_image_checker.py               | 63 +++++++++++++++++++
 2 files changed, 118 insertions(+)
 create mode 100755 checkbox-support/checkbox_support/image_checker.py
 create mode 100644 checkbox-support/checkbox_support/tests/test_image_checker.py

diff --git a/checkbox-support/checkbox_support/image_checker.py b/checkbox-support/checkbox_support/image_checker.py
new file mode 100755
index 0000000000..a571e033ec
--- /dev/null
+++ b/checkbox-support/checkbox_support/image_checker.py
@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+# Copyright 2022 Canonical Ltd.
+# All rights reserved.
+#
+# Written by:
+#   Patrick Chang <patrick.chang@canonical.com>
+
+import argparse
+from os.path import exists
+from checkbox_support.snap_utils.system import on_ubuntucore
+
+
+def get_type():
+    """
+    Return the type of image.
+    """
+    return "core" if on_ubuntucore() else "classic"
+
+
+def get_source():
+    """
+    Return the source of image.
+    """
+    is_oem_source = False
+
+    if get_type() == "core":
+        try:
+            with open("/run/mnt/ubuntu-seed/.disk/info") as file:
+                lines = [line.rstrip() for line in file]
+                # Only one timestamp such as 20221021.4 if it's stock image
+                # There're three lines in info file if it's oem image
+                is_oem_source = (len(lines) > 1)
+        except FileNotFoundError as e:
+            print(e)
+            return 'unknown'
+    else:
+        is_oem_source = exists("/var/lib/ubuntu_dist_channel")
+
+    return 'oem' if is_oem_source else 'stock'
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-t', '--get_type', action='store_true')
+    parser.add_argument('-s', '--get_source', action='store_true')
+    args = parser.parse_args()
+
+    if args.get_type:
+        print("type: {}".format(get_type()))
+    if args.get_source:
+        print("source: {}".format(get_source()))
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/checkbox-support/checkbox_support/tests/test_image_checker.py b/checkbox-support/checkbox_support/tests/test_image_checker.py
new file mode 100644
index 0000000000..0f69432fac
--- /dev/null
+++ b/checkbox-support/checkbox_support/tests/test_image_checker.py
@@ -0,0 +1,63 @@
+# This file is part of Checkbox.
+#
+# Copyright 2024 Canonical Ltd.
+# Written by:
+#   Vincent Liao <vincent.liao@canonical.com>
+#
+# Checkbox is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License version 3,
+# as published by the Free Software Foundation.
+#
+# Checkbox is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with Checkbox.  If not, see <http://www.gnu.org/licenses/>.
+
+import unittest
+from unittest.mock import patch, mock_open
+from checkbox_support.image_checker import get_type, get_source, main
+from io import StringIO
+
+
+class ImageCheckerTest(unittest.TestCase):
+    mock_info = ("Sun, 20 June 1999 66:66:66 +7777\n"
+                 "Somethingcool\n"
+                 "iot-vincent-core-24-x24-2023-10-31-24\n")
+
+    @patch("checkbox_support.image_checker.on_ubuntucore")
+    def test_get_type(self, mock_on_ubuntu_core):
+        mock_on_ubuntu_core.return_value = True
+        self.assertEqual(get_type(), "core")
+
+        mock_on_ubuntu_core.return_value = False
+        self.assertEqual(get_type(), "classic")
+
+    @patch("checkbox_support.image_checker.exists")
+    @patch("checkbox_support.image_checker.get_type")
+    def test_get_source(self, mock_get_type, mock_exists):
+        # Test when it is core image
+        mock_get_type.return_value = "core"
+        with patch("builtins.open", mock_open(read_data=self.mock_info)):
+            self.assertEqual(get_source(), "oem")
+
+        with patch("builtins.open", side_effect=FileNotFoundError):
+            self.assertEqual(get_source(), "unknown")
+
+        # Test when it is classic image
+        mock_get_type.return_value = "classic"
+        mock_exists.return_value = True
+        self.assertEqual(get_source(), "oem")
+
+        mock_exists.return_value = False
+        self.assertEqual(get_source(), "stock")
+
+    @patch("sys.argv", ["script_name.py", "--get_type", "--get_source"])
+    @patch("checkbox_support.image_checker.get_source")
+    @patch("checkbox_support.image_checker.get_type")
+    def test_main(self, mock_get_type, mock_get_source):
+        main()
+        mock_get_type.assert_called_with()
+        mock_get_source.assert_called_with()