composer-cli: Add a get_arg function

This is in preperation for adding more optional arguments. Adds a
generic get_arg function, tests for it, and converts get_size to use it.

(cherry picked from commit 4a4128af23)

Related: rhbz#1859680

src/composer/cli/compose.py

@@ -24,7 +24,7 @@ import pytoml as toml
 
 from composer import http_client as client
 from composer.cli.help import compose_help
-from composer.cli.utilities import argify, handle_api_result, packageNEVRA
+from composer.cli.utilities import argify, handle_api_result, packageNEVRA, get_arg
 
 def compose_cmd(opts):
     """Process compose commands for API v0
@@ -93,8 +93,8 @@ def run_command(opts, cmd_map):
 def get_size(args):
     """Return optional size argument, and remaining args
 
-    :param api: Details about the API server, "version" and "backend"
-    :type api: dict
+    :param args: list of arguments
+    :type args: list of strings
     :returns: (args, size)
     :rtype: tuple
 
@@ -105,20 +105,10 @@ def get_size(args):
     - multiply by 1024**2 to make it easier on users to specify large sizes
 
     """
-    if len(args) == 0:
-        return (args, 0)
+    args, value = get_arg(args, "--size", int)
+    value = value * 1024**2 if value is not None else 0
+    return (args, value)
 
-    if args[0] != "--size" and "--size" in args[1:]:
-        raise RuntimeError("--size must be first argument after the command")
-    if args[0] != "--size":
-        return (args, 0)
-
-    if len(args) < 2:
-        raise RuntimeError("--size is missing the value, in MiB")
-
-    # Let this raise an error for non-digit input
-    size = int(args[1])
-    return (args[2:], size * 1024**2)
 
 def compose_list(socket_path, api_version, args, show_json=False, testmode=0):
     """Return a simple list of compose identifiers"""

src/composer/cli/help.py

@@ -20,9 +20,9 @@ compose start [--size XXXX] <BLUEPRINT> <TYPE> [<IMAGE-NAME> <PROVIDER> <PROFILE
     Start a compose using the selected blueprint and output type. --size is in MiB.
     NOTE: uploading and --size are only supported by osbuild-composer.
 
-compose start-ostree [--size XXXX] <BLUEPRINT> <TYPE> <REF> <PARENT> [<IMAGE-NAME> <PROVIDER> <PROFILE> | <IMAGE-NAME> <PROFILE.TOML>]
+compose start-ostree [--size XXXX] [--parent PARENT] [--ref REF] <BLUEPRINT> <TYPE> [<IMAGE-NAME> <PROFILE.TOML>]
     Start an ostree compose using the selected blueprint and output type. Optionally start an upload. This command
-    is only supported by osbuild-composer, and requires the ostree REF and PARENT. --size is in MiB.
+    is only supported by osbuild-composer. --size is in MiB.
 
 compose types
     List the supported output types.

src/composer/cli/utilities.py

@@ -93,3 +93,31 @@ def packageNEVRA(pkg):
         return "%s-%s:%s-%s.%s" % (pkg["name"], pkg["epoch"], pkg["version"], pkg["release"], pkg["arch"])
     else:
         return "%s-%s-%s.%s" % (pkg["name"], pkg["version"], pkg["release"], pkg["arch"])
+
+def get_arg(args, name, argtype=None):
+    """Return optional value from args, and remaining args
+
+    :param args: list of arguments
+    :type args: list of strings
+    :param name: The argument to remove from the args list
+    :type name: string
+    :param argtype: Type to use for checking the argument value
+    :type argtype: type
+    :returns (args, value)
+    :rtype: tuple
+
+    This removes the optional argument and value from the argument list, returns the new list,
+    and the value of the argument.
+    """
+    try:
+        idx = args.index(name)
+        if len(args) < idx+2:
+            raise RuntimeError(f"{name} is missing the value")
+        value = args[idx+1]
+    except ValueError:
+        return (args, None)
+
+    if argtype:
+        value = argtype(value)
+
+    return (args[:idx]+args[idx+2:], value)

tests/composer/test_compose.py

@@ -382,8 +382,7 @@ class SizeTest(unittest.TestCase):
                          (["blueprint", "type", "imagename", "profile"], 0))
 
     def test_size_later(self):
-        with self.assertRaises(RuntimeError):
-            get_size(["blueprint", "--size", "100", "type"])
+        self.assertEqual(get_size(["start", "--size", "100", "type"]), (["start", "type"], 104857600))
 
     def test_size_no_value(self):
         with self.assertRaises(RuntimeError):

tests/composer/test_utilities.py

@@ -18,7 +18,7 @@ import unittest
 
 from pylorax.api.errors import INVALID_CHARS
 from composer.cli.utilities import argify, toml_filename, frozen_toml_filename, packageNEVRA
-from composer.cli.utilities import handle_api_result
+from composer.cli.utilities import handle_api_result, get_arg
 
 class CliUtilitiesTest(unittest.TestCase):
     def test_argify(self):
@@ -94,3 +94,37 @@ class CliUtilitiesTest(unittest.TestCase):
         """Test a result with show_json=True, errors=[], and no status field"""
         result = {"foo": "bar", "errors": []}
         self.assertEqual(handle_api_result(result, show_json=True), (0, True))
+
+    def test_get_arg_empty(self):
+        """Test get_arg with no arguments"""
+        self.assertEqual(get_arg([], "--size"), ([], None))
+
+    def test_get_arg_no_arg(self):
+        """Test get_arg with no argument in the list"""
+        self.assertEqual(get_arg(["first", "second"], "--size"), (["first", "second"], None))
+
+    def test_get_arg_notype(self):
+        """Test get_arg with no argtype set"""
+        self.assertEqual(get_arg(["first", "--size", "100", "second"], "--size"), (["first", "second"], "100"))
+
+    def test_get_arg_string(self):
+        """Test get_arg with a string argument"""
+        self.assertEqual(get_arg(["first", "--size", "100", "second"], "--size", str), (["first", "second"], "100"))
+
+    def test_get_arg_int(self):
+        """Test get_arg with an int argument"""
+        self.assertEqual(get_arg(["first", "--size", "100", "second"], "--size", int), (["first", "second"], 100))
+
+    def test_get_arg_short(self):
+        """Test get_arg error handling with a short list"""
+        with self.assertRaises(RuntimeError):
+            get_arg(["first", "--size", ], "--size", int)
+
+    def test_get_arg_start(self):
+        """Test get_arg with the argument at the start of the list"""
+        self.assertEqual(get_arg(["--size", "100", "first", "second"], "--size", int), (["first", "second"], 100))
+
+    def test_get_arg_wrong_type(self):
+        """Test get_arg with the wrong type"""
+        with self.assertRaises(ValueError):
+            get_arg(["first", "--size", "abc", "second"], "--size", int)