Add test_image_checker in to checkbox_support to check the image (New)

checkbox-support/checkbox_support/scripts/image_checker.py

@@ -0,0 +1,54 @@
+# 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 = file.readlines()
+                # 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", "--type", action="store_true")
+    parser.add_argument("-s", "--source", action="store_true")
+    args = parser.parse_args()
+
+    if args.type:
+        print("type: {}".format(get_type()))
+    if args.source:
+        print("source: {}".format(get_source()))
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

checkbox-support/checkbox_support/tests/test_image_checker.py

@@ -0,0 +1,68 @@
+# 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.scripts.image_checker import get_type, get_source, main
+from io import StringIO
+
+
+class ImageCheckerTest(unittest.TestCase):
+
+    @patch("checkbox_support.scripts.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.scripts.image_checker.exists")
+    @patch("checkbox_support.scripts.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"
+        mock_oem_info = ("Sun, 20 June 1999 66:66:66 +7777\n"
+                         "Somethingcool\n"
+                         "iot-vincent-core-24-x24-2023-10-31-24\n")
+        with patch("builtins.open", mock_open(read_data=mock_oem_info)):
+            self.assertEqual(get_source(), "oem")
+
+        mock_stock_info = "20231031.6"
+        with patch("builtins.open", mock_open(read_data=mock_stock_info)):
+            self.assertEqual(get_source(), "stock")
+
+        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", "--type", "--source"])
+    @patch("checkbox_support.scripts.image_checker.get_source")
+    @patch("checkbox_support.scripts.image_checker.get_type")
+    @patch("sys.stdout", new_callable=StringIO)
+    def test_main(self, mock_stdout, mock_get_type, mock_get_source):
+        main()
+        mock_get_type.assert_called_with()
+        mock_get_source.assert_called_with()

checkbox-support/pyproject.toml

@@ -34,6 +34,7 @@
   checkbox-support-lsusb = "checkbox_support.scripts.lsusb:main"
   checkbox-support-parse = "checkbox_support.parsers:main"
   checkbox-support-zapper-proxy = "checkbox_support.scripts.zapper_proxy:main"
+  checkbox-support-image_checker = "checkbox_support.scripts.image_checker:main"
 [project.entry-points."plainbox.parsers"]
   pactl-list="checkbox_support.parsers.pactl:parse_pactl_output"
   udevadm="checkbox_support.parsers.udevadm:parse_udevadm_output"

providers/resource/jobs/resource.pxu

@@ -476,3 +476,12 @@ _summary: Collect information about the audio card
 _description: Gets audio resource info from /proc/asound/pcm
 command: audio_card_resource.py
 
+id: image_source_and_type
+estimated_duration: 0.05
+plugin: resource
+_summary: Collect the source and type of image
+_description:
+    Get the source and type of image.
+    Source: 'stock' or 'oem'
+    Type: 'classic' or 'core'
+command: checkbox-support-image_checker -t -s