repo: allow for pre_enable/pre_disable messaging interactions

Consolidate messging hook processing under a single function
handle_message_operations

This is groundwork for FIPS pre-enable and pre-disable custom
messaging and prompts for #1031.

uaclient/entitlements/repo.py

@@ -4,7 +4,16 @@
 import re
 
 try:
-    from typing import Any, Dict, List, Optional, Tuple, Union  # noqa: F401
+    from typing import (  # noqa: F401
+        Any,
+        Callable,
+        Dict,
+        List,
+        Optional,
+        Sequence,
+        Tuple,
+        Union,
+    )
 except ImportError:
     # typing isn't available on trusty, so ignore its absence
     pass
@@ -38,9 +47,9 @@ def repo_pin_priority(self) -> "Union[int, str, None]":
     def disable_apt_auth_only(self) -> bool:
         return False  # Set True on ESM to only remove apt auth
 
-    # Any custom messages to emit pre or post enable or disable operations;
-    # currently post_enable is used in CommonCriteria
-    messaging = {}  # type: Dict[str, List[str]]
+    # Any custom messages to emit to the console or callables which are
+    # handled at pre_enable, pre_disable, pre_install or post_enable stages
+    messaging = {}  # type: Dict[str, List[Union[str, Tuple[Callable, Dict]]]]
 
     @property
     def packages(self) -> "List[str]":
@@ -64,12 +73,16 @@ def enable(self, *, silent_if_inapplicable: bool = False) -> bool:
         """
         if not self.can_enable(silent=silent_if_inapplicable):
             return False
+        msg_ops = self.messaging.get("pre_enable", [])
+        if not handle_message_operations(msg_ops):
+            return False
         self.setup_apt_config()
         if self.packages:
             try:
                 print("Installing {title} packages".format(title=self.title))
-                for msg in self.messaging.get("pre_install", []):
-                    print(msg)
+                msg_ops = self.messaging.get("pre_install", [])
+                if not handle_message_operations(msg_ops):
+                    return False
                 apt.run_apt_command(
                     ["apt-get", "install", "--assume-yes"] + self.packages,
                     status.MESSAGE_ENABLED_FAILED_TMPL.format(
@@ -80,13 +93,17 @@ def enable(self, *, silent_if_inapplicable: bool = False) -> bool:
                 self._cleanup()
                 raise
         print(status.MESSAGE_ENABLED_TMPL.format(title=self.title))
-        for msg in self.messaging.get("post_enable", []):
-            print(msg)
+        msg_ops = self.messaging.get("post_enable", [])
+        if not handle_message_operations(msg_ops):
+            return False
         return True
 
     def disable(self, silent=False):
         if not self.can_disable(silent):
             return False
+        msg_ops = self.messaging.get("pre_disable", [])
+        if not handle_message_operations(msg_ops):
+            return False
         self._cleanup()
         return True
 
@@ -301,3 +318,25 @@ def remove_apt_config(self):
         apt.run_apt_command(
             ["apt-get", "update"], status.MESSAGE_APT_UPDATE_FAILED
         )
+
+
+def handle_message_operations(
+    msg_ops: "List[Union[str, Tuple[Callable, Dict]]]"
+) -> bool:
+    """Emit messages to the console for user interaction
+
+    :param msg_op: A list of strings or tuples. Any string items are printed.
+        Any tuples will contain a callable and a dict of args to pass to the
+        callable. Callables are expected to return True on success and
+        False upon failure.
+
+    :return: True upon success, False on failure.
+    """
+    for msg_op in msg_ops:
+        if isinstance(msg_op, str):
+            print(msg_op)
+        else:  # Then we are a callable and dict of args
+            functor, args = msg_op
+            if not functor(**args):
+                return False
+    return True

uaclient/entitlements/tests/test_repo.py

@@ -7,7 +7,10 @@
 
 from uaclient import apt
 from uaclient import config
-from uaclient.entitlements.repo import RepoEntitlement
+from uaclient.entitlements.repo import (
+    RepoEntitlement,
+    handle_message_operations,
+)
 from uaclient.entitlements.tests.conftest import machine_token
 from uaclient import exceptions
 from uaclient import status
@@ -318,6 +321,70 @@ def test_enable_passes_silent_if_inapplicable_through(
         expected_call = mock.call(silent=bool(silent_if_inapplicable))
         assert [expected_call] == m_can_enable.call_args_list
 
+    @pytest.mark.parametrize(
+        "pre_enable_msg, output, setup_apt_call_count",
+        (
+            (["msg1", (lambda: False, {}), "msg2"], "msg1\n", 0),
+            (
+                ["msg1", (lambda: True, {}), "msg2"],
+                "msg1\nmsg2\nRepo Test Class enabled\n",
+                1,
+            ),
+        ),
+    )
+    @mock.patch.object(RepoTestEntitlement, "setup_apt_config")
+    @mock.patch.object(RepoTestEntitlement, "can_enable", return_value=True)
+    def test_enable_can_exit_on_pre_enable_messaging_hooks(
+        self,
+        _can_enable,
+        setup_apt_config,
+        pre_enable_msg,
+        output,
+        setup_apt_call_count,
+        entitlement,
+        capsys,
+    ):
+        messaging = {"pre_enable": pre_enable_msg}
+        with mock.patch.object(type(entitlement), "messaging", messaging):
+            with mock.patch.object(type(entitlement), "packages", []):
+                entitlement.enable()
+        stdout, _ = capsys.readouterr()
+        assert output == stdout
+        assert setup_apt_call_count == setup_apt_config.call_count
+
+    @pytest.mark.parametrize(
+        "pre_disable_msg, output, remove_apt_call_count",
+        (
+            (["msg1", (lambda: False, {}), "msg2"], "msg1\n", 0),
+            (["msg1", (lambda: True, {}), "msg2"], "msg1\nmsg2\n", 1),
+        ),
+    )
+    @mock.patch(M_PATH + "util.subp", return_value=("", ""))
+    @mock.patch.object(RepoTestEntitlement, "remove_apt_config")
+    @mock.patch.object(RepoTestEntitlement, "can_disable", return_value=True)
+    def test_enable_can_exit_on_pre_disable_messaging_hooks(
+        self,
+        _can_disable,
+        remove_apt_config,
+        m_subp,
+        pre_disable_msg,
+        output,
+        remove_apt_call_count,
+        entitlement,
+        capsys,
+    ):
+        messaging = {"pre_disable": pre_disable_msg}
+        with mock.patch.object(type(entitlement), "messaging", messaging):
+            with mock.patch.object(type(entitlement), "packages", []):
+                entitlement.disable()
+        stdout, _ = capsys.readouterr()
+        assert output == stdout
+        assert remove_apt_call_count == remove_apt_config.call_count
+        if remove_apt_call_count > 0:
+            assert [
+                mock.call(["apt-get", "remove", "--assume-yes"])
+            ] == m_subp.call_args_list
+
     @pytest.mark.parametrize("with_pre_install_msg", (False, True))
     @pytest.mark.parametrize("packages", (["a"], [], None))
     @mock.patch(M_PATH + "util.subp", return_value=("", ""))
@@ -792,3 +859,34 @@ def test_enabled_status_by_apt_policy(
             expected_explanation = "Repo Test Class is not configured"
         assert expected_status == application_status
         assert expected_explanation == explanation
+
+
+def success_call():
+    print("success")
+    return True
+
+
+def fail_call(a=None):
+    print("fail %s" % a)
+    return False
+
+
+class TestHandleMessageOperations:
+    @pytest.mark.parametrize(
+        "msg_ops, retval, output",
+        (
+            ([], True, ""),
+            (["msg1", "msg2"], True, "msg1\nmsg2\n"),
+            (
+                [(success_call, {}), "msg1", (fail_call, {"a": 1}), "msg2"],
+                False,
+                "success\nmsg1\nfail 1\n",
+            ),
+        ),
+    )
+    def test_handle_message_operations_for_strings_and_callables(
+        self, msg_ops, retval, output, capsys
+    ):
+        assert retval is handle_message_operations(msg_ops)
+        out, _err = capsys.readouterr()
+        assert output == out