util: prompt_for_confirmation custom msg and assume_yes

uaclient/cli.py

@@ -326,7 +326,7 @@ def _detach(cfg: config.UAConfig, assume_yes: bool) -> int:
         print("Detach will disable the following service{}:".format(suffix))
         for ent in to_disable:
             print("    {}".format(ent.name))
-    if not assume_yes and not util.prompt_for_confirmation():
+    if not util.prompt_for_confirmation(assume_yes=assume_yes):
         return 1
     for ent in to_disable:
         ent.disable(silent=True)

uaclient/tests/test_cli_detach.py

@@ -40,7 +40,7 @@ def test_unattached_error_message(self, m_getuid, _m_prompt, FakeConfig):
 
     @pytest.mark.parametrize(
         "prompt_response,assume_yes,expect_disable",
-        [(True, False, True), (False, False, False), (None, True, True)],
+        [(True, False, True), (False, False, False), (True, True, True)],
     )
     @mock.patch("uaclient.cli.entitlements")
     @mock.patch("uaclient.contract.UAContractClient")
@@ -56,8 +56,7 @@ def test_entitlements_disabled_appropriately(
         FakeConfig,
     ):
         # The three parameters:
-        #   prompt_response: the user's response to the prompt, or None if no
-        #                    prompt should be displayed
+        #   prompt_response: the user's response to the prompt
         #   assume_yes: the value of the --assume-yes flag in the args passed
         #               to the action
         #   expect_disable: whether or not the enabled entitlement is expected
@@ -68,10 +67,7 @@ def test_entitlements_disabled_appropriately(
         fake_client = FakeContractClient(cfg)
         m_client.return_value = fake_client
 
-        if prompt_response is not None:
-            m_prompt.return_value = prompt_response
-        else:
-            m_prompt.side_effect = Exception("SHOULD NOT BE CALLED")
+        m_prompt.return_value = prompt_response
 
         m_entitlements.ENTITLEMENT_CLASSES = [
             entitlement_cls_mock_factory(False),
@@ -103,6 +99,7 @@ def test_entitlements_disabled_appropriately(
         else:
             assert 0 == disabled_cls.return_value.disable.call_count
             assert 1 == return_code
+        assert [mock.call(assume_yes=assume_yes)] == m_prompt.call_args_list
 
     @mock.patch("uaclient.cli.entitlements")
     @mock.patch("uaclient.contract.UAContractClient")

uaclient/tests/test_util.py

@@ -546,7 +546,15 @@ def test_input_conversion(self, m_input, return_value, user_input):
         m_input.return_value = user_input
         assert return_value == util.prompt_for_confirmation()
 
-    def test_prompt_text(self, m_input):
-        util.prompt_for_confirmation()
+    @pytest.mark.parametrize(
+        "assume_yes,message,input_calls",
+        [
+            (True, "message ignored on assume_yes=True", []),
+            (False, "", [mock.call("Are you sure? (y/N) ")]),
+            (False, "Custom yep? (y/N) ", [mock.call("Custom yep? (y/N) ")]),
+        ]
+    )
+    def test_prompt_text(self, m_input, assume_yes, message, input_calls):
+        util.prompt_for_confirmation(msg=message, assume_yes=assume_yes)
 
-        assert [mock.call("Are you sure? (y/N) ")] == m_input.call_args_list
+        assert input_calls == m_input.call_args_list

uaclient/util.py

@@ -387,14 +387,22 @@ def parse_os_release(release_file: "Optional[str]" = None) -> "Dict[str, str]":
     return data
 
 
-def prompt_for_confirmation() -> bool:
+def prompt_for_confirmation(msg: str = "", assume_yes: bool = False) -> bool:
     """
     Display a confirmation prompt, returning a bool indicating the response
 
+    :param msg: String custom prompt text to emit from input call.
+    :param assume_yes: Boolean set True to skip confirmation input and return
+        True.
+
     This function will only prompt a single time, and defaults to "no" (i.e. it
     returns False).
     """
-    value = input("Are you sure? (y/N) ")
+    if assume_yes:
+        return True
+    if not msg:
+        msg = "Are you sure? (y/N) "
+    value = input(msg)
     if value.lower().strip() in ["y", "yes"]:
         return True
     return False