Pr1/increase_RebootManager_and_PatchAssessor_code coverage

src/core/src/core_logic/RebootManager.py

@@ -15,8 +15,6 @@
 # Requires Python 2.7+
 
 """Reboot management"""
-import datetime
-import subprocess
 import time
 from core.src.bootstrap.Constants import Constants
 

src/core/tests/Test_PatchAssessor.py

@@ -23,7 +23,7 @@
 from core.src.service_interfaces.TelemetryWriter import TelemetryWriter
 from core.tests.library.ArgumentComposer import ArgumentComposer
 from core.tests.library.RuntimeCompositor import RuntimeCompositor
-from core.src.core_logic.Stopwatch import Stopwatch
+
 
 class TestPatchAssessor(unittest.TestCase):
     def setUp(self):
@@ -165,7 +165,6 @@ def test_write_assessment_perf_logs(self):
         err_msg = "{0}=".format(str(Constants.PerfLogTrackerParams.ERROR_MSG))
         self.assertTrue(err_msg in str(self.runtime.patch_assessor.stopwatch.task_details))
 
-
     def test_stopwatch_properties_assessment_fail(self):
         self.runtime.set_legacy_test_type('UnalignedPath')
         self.assertRaises(Exception, self.runtime.patch_assessor.start_assessment)
@@ -181,6 +180,15 @@ def test_raise_if_min_python_version_not_met(self):
             self.runtime.patch_assessor.start_assessment()
         self.assertEqual(str(context.exception), Constants.PYTHON_NOT_COMPATIBLE_ERROR_MSG.format(sys.version_info))
 
+    def test_patch_assessment_throws_exception(self):
+        self.runtime.package_manager.get_all_updates = lambda: self.raise_ex()
+
+        with self.assertRaises(Exception) as context:
+            self.runtime.patch_assessor.start_assessment()
+
+        self.assertIn(Constants.ERROR_ADDED_TO_STATUS, repr(context.exception))
+        self.assertEqual(context.exception.args[1], "[{0}]".format(Constants.ERROR_ADDED_TO_STATUS))
+
     def raise_ex(self):
         raise Exception()
 

src/core/tests/Test_RebootManager.py

@@ -13,8 +13,8 @@
 # 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
@@ -113,6 +113,34 @@ def test_reboot_always_time_not_available(self):
         self.assertEqual(reboot_manager.start_reboot_if_required_and_time_available(10), False)
         runtime.stop()
 
+    def test_reboot_if_required_no_reboot_pending(self):
+        reboot_setting_in_api = 'IfRequired'
+        argument_composer = ArgumentComposer()
+        argument_composer.reboot_setting = reboot_setting_in_api
+        runtime = RuntimeCompositor(argument_composer.get_composed_arguments(), True, Constants.YUM)
+        reboot_manager = runtime.reboot_manager
+
+        # Validate single reboot scenario
+        runtime.status_handler.is_reboot_pending = False
+        self.assertEqual(reboot_manager.start_reboot_if_required_and_time_available(20), False)
+        runtime.stop()
+
+    def test_start_reboot_raise_exception(self):
+        reboot_setting_in_api = 'Always'
+        argument_composer = ArgumentComposer()
+        argument_composer.reboot_setting = reboot_setting_in_api
+        runtime = RuntimeCompositor(argument_composer.get_composed_arguments(), True, Constants.YUM)
+        Constants.REBOOT_WAIT_TIMEOUT_IN_MINUTES = -20
+
+        with self.assertRaises(Exception) as context:
+            runtime.use_original_rm_start_reboot()
+            runtime.reboot_manager.start_reboot()
+
+        # assert
+        self.assertIn("Reboot failed to proceed on the machine in a timely manner.", repr(context.exception))
+        self.assertEqual(context.exception.args[1], "[{0}]".format(Constants.ERROR_ADDED_TO_STATUS))
+        runtime.stop()
+
 
 if __name__ == '__main__':
     unittest.main()

src/core/tests/library/RuntimeCompositor.py

@@ -43,6 +43,7 @@
 class RuntimeCompositor(object):
     def __init__(self, argv=Constants.DEFAULT_UNSPECIFIED_VALUE, legacy_mode=False, package_manager_name=Constants.APT, vm_cloud_type=Constants.VMCloudType.AZURE):
         # Init data
+        self.original_rm_start_reboot = None
         self.current_env = Constants.DEV
         os.environ[Constants.LPE_ENV_VARIABLE] = self.current_env
         self.argv = argv if argv != Constants.DEFAULT_UNSPECIFIED_VALUE else ArgumentComposer().get_composed_arguments()
@@ -150,11 +151,18 @@ def reconfigure_env_layer_to_legacy_mode(self):
             self.env_layer.etc_environment_file_path = os.getcwd()
 
     def reconfigure_reboot_manager(self):
+        # Preserve the original reboot manager start_reboot method
+        self.original_rm_start_reboot = self.reboot_manager.start_reboot
+
+        # Reassign start_reboot to a new mock method
         self.reboot_manager.start_reboot = self.start_reboot
 
     def start_reboot(self, message="Test initiated reboot mock"):
         self.status_handler.set_installation_reboot_status(Constants.RebootStatus.STARTED)
 
+    def use_original_rm_start_reboot(self):
+        self.reboot_manager.start_reboot = self.original_rm_start_reboot
+
     def reconfigure_package_manager(self):
         self.backup_get_current_auto_os_patch_state = self.package_manager.get_current_auto_os_patch_state
         self.package_manager.get_current_auto_os_patch_state = self.get_current_auto_os_patch_state