Add retry to check_sudo_status and unit test

src/core/src/bootstrap/Bootstrapper.py

@@ -19,6 +19,7 @@
 import json
 import os
 import sys
+import time
 from core.src.bootstrap.ConfigurationFactory import ConfigurationFactory
 from core.src.bootstrap.Constants import Constants
 from core.src.bootstrap.Container import Container
@@ -142,10 +143,29 @@ def basic_environment_health_check(self):
         self.composite_logger.log("Process id: " + str(os.getpid()))
 
         # Ensure sudo works in the environment
-        sudo_check_result = self.check_sudo_status()
+        sudo_check_result = self.check_sudo_status_with_retry()
         self.composite_logger.log_debug("Sudo status check: " + str(sudo_check_result) + "\n")
 
+    def check_sudo_status_with_retry(self, raise_if_not_sudo=True):
+        # type:(bool) -> bool
+        """ retry to invoke sudo check """
+        for attempts in range(1, Constants.MAX_CHECK_SUDO_RETRY_COUNT + 1):
+            try:
+                sudo_status = self.check_sudo_status(raise_if_not_sudo=raise_if_not_sudo)
+                if sudo_status and attempts > 1:
+                    self.composite_logger.log_debug("Sudo Check Successfully [RetryCount={0}][MaxRetryCount={1}]".format(str(attempts), Constants.MAX_CHECK_SUDO_RETRY_COUNT))
+                return sudo_status
+            except Exception as exception:
+                if attempts >= Constants.MAX_CHECK_SUDO_RETRY_COUNT:
+                    self.composite_logger.log_error("Customer environment error (sudo failure). [Exception={0}][RetryCount={1}]".format(str(exception), str(attempts)))
+                    if raise_if_not_sudo:
+                        raise
+
+                self.composite_logger.log_debug("Retrying sudo status check after a delay of [ElapsedTimeInSeconds={0}][RetryCount={1}]".format(Constants.MAX_CHECK_SUDO_INTERVAL_IN_SEC, str(attempts)))
+                time.sleep(Constants.MAX_CHECK_SUDO_INTERVAL_IN_SEC)
+
     def check_sudo_status(self, raise_if_not_sudo=True):
+        # type:(bool) -> bool
         """ Checks if we can invoke sudo successfully. """
         try:
             self.composite_logger.log("Performing sudo status check... This should complete within 10 seconds.")
@@ -170,7 +190,7 @@ def check_sudo_status(self, raise_if_not_sudo=True):
             else:
                 raise Exception("Unexpected sudo check result. Output: " + " ".join(output.split("\n")))
         except Exception as exception:
-            self.composite_logger.log_error("Sudo status check failed. Please ensure the computer is configured correctly for sudo invocation. " +
+            self.composite_logger.log_debug("Sudo status check failed. Please ensure the computer is configured correctly for sudo invocation. " +
                                             "Exception details: " + str(exception))
             if raise_if_not_sudo:
                 raise

src/core/src/bootstrap/Constants.py

@@ -220,6 +220,9 @@ class StatusTruncationConfig(EnumBackport):
     MAX_IMDS_CONNECTION_RETRY_COUNT = 5
     MAX_ZYPPER_REPO_REFRESH_RETRY_COUNT = 5
     MAX_COMPLETE_STATUS_FILES_TO_RETAIN = 10
+    SET_CHECK_SUDO_STATUS_TRUE = True
+    MAX_CHECK_SUDO_RETRY_COUNT = 6
+    MAX_CHECK_SUDO_INTERVAL_IN_SEC = 300
 
     class PackageBatchConfig(EnumBackport):
         # Batch Patching Parameters

src/core/tests/Test_Bootstrapper.py

@@ -0,0 +1,116 @@
+# Copyright 2020 Microsoft Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Requires Python 2.7+
+import unittest
+
+from core.src.bootstrap.Constants import Constants
+from core.tests.library.ArgumentComposer import ArgumentComposer
+from core.tests.library.RuntimeCompositor import RuntimeCompositor
+
+
+class TestBootstrapper(unittest.TestCase):
+    def setUp(self):
+        self.sudo_check_status_attempts = 0
+        Constants.SET_CHECK_SUDO_STATUS_TRUE = False
+        argument_composer = ArgumentComposer()
+        argument_composer.operation = Constants.ASSESSMENT
+        self.argv = argument_composer.get_composed_arguments()
+        self.runtime = RuntimeCompositor(self.argv, legacy_mode=True, package_manager_name=Constants.APT)
+
+    def tearDown(self):
+        self.sudo_check_status_attempts = 0
+        Constants.SET_CHECK_SUDO_STATUS_TRUE = True
+        self.runtime.stop()
+
+    # regions mock
+    def mock_false_run_command_output(self, command, no_output=False, chk_err=True):
+        """Mock a failed sudo check status command output to test retry logic."""
+        # Mock failure to trigger retry logic in check_sudo_status
+        return (1, "[sudo] password for user:\nFalse")
+
+    def mock_insufficient_run_command_output(self, command, no_output=False, chk_err=True):
+        """Mock a insufficient output line in sudo check status command output."""
+        # Mock failure to trigger retry logic in check_sudo_status
+        return (1, "[sudo] password for user:")
+
+    def mock_unexpected_output_run_command_output(self, command, no_output=False, chk_err=True):
+        """Mock a insufficient output line in sudo check status command output."""
+        # Mock failure to trigger retry logic in check_sudo_status
+        return (1, "[sudo] password for user:\nUnexpectedOutput")
+
+    def mock_retry_run_command_output(self, command, no_output=False, chk_err=True):
+        """Mock 3 failed sudo check status attempts followed by a success on the 4th attempt."""
+        self.sudo_check_status_attempts += 1
+
+        # Mock failure on the first two attempts
+        if self.sudo_check_status_attempts <= 2:
+            return (1, "[sudo] password for user:\nFalse")
+
+        # Mock success (True) on the 3rd attempt
+        elif self.sudo_check_status_attempts == 3:
+            return (0, "uid=0(root) gid=0(root) groups=0(root)\nTrue")
+    # end regions mock
+
+    def test_check_sudo_status_all_attempts_failed(self):
+        # Set raise_if_not_sudo=False to test the `return False` all attempts failed
+        self.runtime.env_layer.run_command_output = self.mock_false_run_command_output
+        result = self.runtime.bootstrapper.check_sudo_status_with_retry(raise_if_not_sudo=False)
+        self.assertFalse(result, "Expected check_sudo_status to return False after all attempts failed")
+
+    def test_check_sudo_status_throw_exception(self):
+        # Set raise_if_not_sudo=True to throw exception) after all retries
+        self.runtime.env_layer.run_command_output = self.mock_false_run_command_output
+        with self.assertRaises(Exception) as context:
+            self.runtime.bootstrapper.check_sudo_status_with_retry(raise_if_not_sudo=True)
+
+        # Verify exception msg contains the expected failure text
+        self.assertTrue("Unable to invoke sudo successfully" in str(context.exception))
+
+    def test_check_sudo_status_insufficient_output_lines(self):
+        # Test insufficient output lines to raise exception after all retries
+        self.runtime.env_layer.run_command_output = self.mock_insufficient_run_command_output
+
+        with self.assertRaises(Exception) as context:
+            self.runtime.bootstrapper.check_sudo_status_with_retry()
+
+        # Verify exception msg contains the expected failure text
+        self.assertTrue("Unexpected sudo check result" in str(context.exception))
+
+    def test_check_sudo_status_unexpected_output_lines(self):
+        # Test unexpected output with neither false or true to raise exception after all retries
+        self.runtime.env_layer.run_command_output = self.mock_unexpected_output_run_command_output
+
+        with self.assertRaises(Exception) as context:
+            self.runtime.bootstrapper.check_sudo_status_with_retry()
+
+        # Verify exception msg contains the expected failure text
+        self.assertTrue("Unexpected sudo check result" in str(context.exception))
+
+    def test_check_sudo_status_succeeds_on_third_attempt(self):
+        # Test retry logic in check sudo status after 2 failed attempts followed by success (true)
+        self.runtime.env_layer.run_command_output = self.mock_retry_run_command_output
+
+        # Attempt to check sudo status, succeed (true) on the 3rd attempt
+        result = self.runtime.bootstrapper.check_sudo_status_with_retry(raise_if_not_sudo=True)
+
+        # Verify the result is success (True)
+        self.assertTrue(result, "Expected check_sudo_status to succeed on the 3rd attempts")
+
+        # Verify 3 attempts were made
+        self.assertEqual(self.sudo_check_status_attempts, 3, "Expected exactly 3 attempts in check_sudo_status")
+
+
+if __name__ == '__main__':
+    unittest.main()

src/core/tests/library/RuntimeCompositor.py

@@ -71,28 +71,32 @@ def mkdtemp_runner():
         urlreq.urlopen = self.mock_urlopen
 
         # Adapted bootstrapper
-        bootstrapper = Bootstrapper(self.argv, capture_stdout=False)
+        self.bootstrapper = Bootstrapper(self.argv, capture_stdout=False)
 
-        # Overriding sudo status check
-        Bootstrapper.check_sudo_status = self.check_sudo_status
+        # Store the original check_sudo_status method to reset when needed in stop()
+        self.original_check_sudo_status = Bootstrapper.check_sudo_status
 
         # Reconfigure env layer for legacy mode tests
-        self.env_layer = bootstrapper.env_layer
+        self.env_layer = self.bootstrapper.env_layer
         if legacy_mode:
             self.legacy_env_layer_extensions = LegacyEnvLayerExtensions(package_manager_name, test_type)
             self.reconfigure_env_layer_to_legacy_mode()
 
+        # Overriding check_sudo_status to always true
+        if Constants.SET_CHECK_SUDO_STATUS_TRUE:
+            Bootstrapper.check_sudo_status = self.check_sudo_status
+
         # Core components
-        self.container = bootstrapper.build_out_container()
-        self.file_logger = bootstrapper.file_logger
-        self.composite_logger = bootstrapper.composite_logger
+        self.container = self.bootstrapper.build_out_container()
+        self.file_logger = self.bootstrapper.file_logger
+        self.composite_logger = self.bootstrapper.composite_logger
 
         # re-initializing telemetry_writer, outside of Bootstrapper, to correctly set the env_layer configured for tests
-        self.telemetry_writer = TelemetryWriter(self.env_layer, self.composite_logger, bootstrapper.telemetry_writer.events_folder_path, bootstrapper.telemetry_supported)
-        bootstrapper.telemetry_writer = self.telemetry_writer
-        bootstrapper.composite_logger.telemetry_writer = self.telemetry_writer
+        self.telemetry_writer = TelemetryWriter(self.env_layer, self.composite_logger, self.bootstrapper.telemetry_writer.events_folder_path, self.bootstrapper.telemetry_supported)
+        self.bootstrapper.telemetry_writer = self.telemetry_writer
+        self.bootstrapper.composite_logger.telemetry_writer = self.telemetry_writer
 
-        self.lifecycle_manager, self.status_handler = bootstrapper.build_core_components(self.container)
+        self.lifecycle_manager, self.status_handler = self.bootstrapper.build_core_components(self.container)
 
         # Business logic components
         self.execution_config = self.container.get('execution_config')
@@ -107,7 +111,7 @@ def mkdtemp_runner():
         self.patch_assessor = self.container.get('patch_assessor')
         self.patch_installer = self.container.get('patch_installer')
         self.maintenance_window = self.container.get('maintenance_window')
-        self.vm_cloud_type = bootstrapper.configuration_factory.vm_cloud_type
+        self.vm_cloud_type = self.bootstrapper.configuration_factory.vm_cloud_type
         # Extension handler dependency
         self.write_ext_state_file(self.lifecycle_manager.ext_state_file_path, self.execution_config.sequence_number, datetime.datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ"), self.execution_config.operation)
 
@@ -127,6 +131,7 @@ def mkdtemp_runner():
     def stop(self):
         self.file_logger.close(message_at_close="<Runtime stopped>")
         self.container.reset()
+        Bootstrapper.check_sudo_status = self.original_check_sudo_status
 
     @staticmethod
     def write_ext_state_file(path, sequence_number, achieve_enable_by, operation):