Add ordered dictionary for quicker look up of Assessment packages in StatusHandler

src/core/src/service_interfaces/StatusHandler.py

@@ -20,6 +20,7 @@
 import shutil
 import time
 from core.src.bootstrap.Constants import Constants
+from collections import OrderedDict
 
 
 class StatusHandler(object):
@@ -53,6 +54,7 @@ def __init__(self, env_layer, execution_config, composite_logger, telemetry_writ
         self.__assessment_packages = []
         self.__assessment_errors = []
         self.__assessment_total_error_count = 0  # All errors during assess, includes errors not in error objects due to size limit
+        self.__assessment_packages_map = OrderedDict()
 
         # Internal in-memory representation of Patch Metadata for HealthStore
         self.__metadata_for_healthstore_substatus_json = None
@@ -101,18 +103,21 @@ def reset_assessment_data(self):
         self.__assessment_packages = []
         self.__assessment_errors = []
         self.__assessment_total_error_count = 0
+        self.__assessment_packages_map = OrderedDict()
 
     def set_package_assessment_status(self, package_names, package_versions, classification="Other", status="Available"):
         """ Externally available method to set assessment status for one or more packages of the **SAME classification and status** """
         self.composite_logger.log_debug("Setting package assessment status in bulk. [Count={0}]".format(str(len(package_names))))
+
         for package_name, package_version in zip(package_names, package_versions):
             patch_already_saved = False
             patch_id = self.__get_patch_id(package_name, package_version)
-            for i in range(0, len(self.__assessment_packages)):
-                if patch_id == self.__assessment_packages[i]['patchId']:
-                    patch_already_saved = True
-                    self.__assessment_packages[i]['classifications'] = [classification]
-                    # self.__assessment_packages[i]['patchState'] = status
+
+            # Match patch_id in map and update existing patch's classification i.e from other -> security
+            if len(self.__assessment_packages_map) > 0 and patch_id in self.__assessment_packages_map:
+                self.__assessment_packages_map.setdefault(patch_id, {})['classifications'] = [classification]
+                # self.__assessment_packages_map.setdefault(patch_id, {})['patchState'] = status
+                patch_already_saved = True
 
             if patch_already_saved is False:
                 record = {
@@ -122,8 +127,10 @@ def set_package_assessment_status(self, package_names, package_versions, classif
                     "classifications": [classification]
                     # "patchState": str(status) # Allows for capturing 'Installed' packages in addition to 'Available', when commented out, if spec changes
                 }
-                self.__assessment_packages.append(record)
+                # Add new patch to map
+                self.__assessment_packages_map[patch_id] = record
 
+        self.__assessment_packages = list(self.__assessment_packages_map.values())
         self.__assessment_packages = self.sort_packages_by_classification_and_state(self.__assessment_packages)
         self.set_assessment_substatus_json()
 
@@ -540,6 +547,7 @@ def load_status_file_components(self, initial_load=False):
         self.__assessment_summary_json = None
         self.__assessment_packages = []
         self.__assessment_errors = []
+        self.__assessment_packages_map = OrderedDict()
 
         self.__metadata_for_healthstore_substatus_json = None
         self.__metadata_for_healthstore_summary_json = None
@@ -601,7 +609,8 @@ def load_status_file_components(self, initial_load=False):
             if name == Constants.PATCH_ASSESSMENT_SUMMARY:     # if it exists, it must be to spec, or an exception will get thrown
                 message = status_file_data['status']['substatus'][i]['formattedMessage']['message']
                 self.__assessment_summary_json = json.loads(message)
-                self.__assessment_packages = self.__assessment_summary_json['patches']
+                self.__assessment_packages_map = OrderedDict((package["patchId"], package) for package in self.__assessment_summary_json['patches'])
+                self.__assessment_packages = list(self.__assessment_packages_map.values())
                 errors = self.__assessment_summary_json['errors']
                 if errors is not None and errors['details'] is not None:
                     self.__assessment_errors = errors['details']

src/core/tests/Test_StatusHandler.py

@@ -396,6 +396,37 @@ def test_sort_packages_by_classification_and_state(self):
             self.assertEqual(installation_patches_sorted[12]["name"], "test-package-2")  # | Other              | Excluded    |
             self.assertEqual(installation_patches_sorted[13]["name"], "test-package-1")  # | Other              | NotSelected |
 
+    def test_assessment_packages_map(self):
+        patch_count_for_test = 5
+        expected_patch_id = 'python-samba0_2:4.4.5+dfsg-2ubuntu5.4_Ubuntu_16.04'
+
+        status_handler = StatusHandler(self.runtime.env_layer, self.runtime.execution_config, self.runtime.composite_logger, self.runtime.telemetry_writer, self.runtime.vm_cloud_type)
+        self.runtime.execution_config.operation = Constants.ASSESSMENT
+        self.runtime.status_handler.set_current_operation(Constants.ASSESSMENT)
+
+        test_packages, test_package_versions = self.__set_up_packages_func(patch_count_for_test)
+        status_handler.set_package_assessment_status(test_packages, test_package_versions, 'Critical')
+
+        with self.runtime.env_layer.file_system.open(self.runtime.execution_config.status_file_path, 'r') as file_handle:
+            substatus_file_data = json.load(file_handle)[0]["status"]["substatus"][0]
+
+        self.assertTrue(substatus_file_data["name"] == Constants.PATCH_ASSESSMENT_SUMMARY)
+        formatted_message = json.loads(substatus_file_data['formattedMessage']['message'])
+        self.assertEqual(len(formatted_message['patches']), patch_count_for_test)
+        self.assertEqual(formatted_message['patches'][0]['classifications'], ['Critical'])
+        self.assertEqual(formatted_message['patches'][0]['name'], 'python-samba0')
+        self.assertEqual(formatted_message['patches'][0]['patchId'], expected_patch_id)
+
+    # Setup functions to populate packages and versions for truncation
+    def __set_up_packages_func(self, val):
+        test_packages = []
+        test_package_versions = []
+
+        for i in range(0, val):
+            test_packages.append('python-samba' + str(i))
+            test_package_versions.append('2:4.4.5+dfsg-2ubuntu5.4')
+
+        return test_packages, test_package_versions
 
 if __name__ == '__main__':
     unittest.main()