From fec8fec434f78b553ee904420c8462368e98f7c2 Mon Sep 17 00:00:00 2001
From: Renan Rodrigo <renanrodrigo@canonical.com>
Date: Fri, 12 Jul 2024 07:09:32 -0300
Subject: [PATCH] cli: add subcommand functionality to commands

Signed-off-by: Renan Rodrigo <renanrodrigo@canonical.com>
---
 uaclient/cli/commands.py                | 12 ++++++++++++
 uaclient/cli/tests/test_cli_commands.py | 23 +++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/uaclient/cli/commands.py b/uaclient/cli/commands.py
index 7f2a8d2df2..440ff10593 100644
--- a/uaclient/cli/commands.py
+++ b/uaclient/cli/commands.py
@@ -1,6 +1,7 @@
 import argparse
 from typing import Callable, Iterable, Optional, Union
 
+from uaclient import messages
 from uaclient.cli.constants import NAME, USAGE_TMPL
 
 
@@ -79,6 +80,7 @@ def __init__(
         action: Callable = lambda *args, **kwargs: None,
         preserve_description: bool = False,
         argument_groups: Iterable[ProArgumentGroup] = (),
+        subcommands: Iterable["ProCommand"] = (),
     ):
         self.name = name
         self.help = help
@@ -87,6 +89,7 @@ def __init__(
         self.action = action
         self.preserve_description = preserve_description
         self.argument_groups = argument_groups
+        self.subcommands = subcommands
 
     def register(self, subparsers: argparse._SubParsersAction):
         self.parser = subparsers.add_parser(
@@ -102,3 +105,12 @@ def register(self, subparsers: argparse._SubParsersAction):
             argument_group.register(self.parser)
 
         self.parser.set_defaults(action=self.action)
+
+        if self.subcommands:
+            subparsers = self.parser.add_subparsers(
+                title=messages.CLI_AVAILABLE_COMMANDS,
+                dest="command",
+                metavar="",
+            )
+            for command in self.subcommands:
+                command.register(subparsers)
diff --git a/uaclient/cli/tests/test_cli_commands.py b/uaclient/cli/tests/test_cli_commands.py
index 5668d50d6f..96e3ca1d23 100644
--- a/uaclient/cli/tests/test_cli_commands.py
+++ b/uaclient/cli/tests/test_cli_commands.py
@@ -66,6 +66,29 @@ def test_command_register(self):
             mock.call(action=example_command.action)
         ] == example_command.parser.set_defaults.call_args_list
 
+    def test_has_subcommands(self):
+        mock_subparsers = mock.MagicMock()
+
+        example_subcommand1 = mock.MagicMock()
+        example_subcommand2 = mock.MagicMock()
+
+        example_command = ProCommand(
+            "example",
+            help="help",
+            description="description",
+            subcommands=[example_subcommand1, example_subcommand2],
+        )
+
+        example_command.register(mock_subparsers)
+
+        inner_subparsers = example_command.parser.add_subparsers.return_value
+        assert [
+            mock.call(inner_subparsers)
+        ] == example_subcommand1.register.call_args_list
+        assert [
+            mock.call(inner_subparsers)
+        ] == example_subcommand2.register.call_args_list
+
 
 class TestProArgument:
     def test_argument_register(self):