fix: merge control markdown with json (#1528) (#1740)

* bug: add test for #1528 Signed-off-by: d10n <[email protected]> * fix: merge control markdown with json (#1528) Function cleanup / unit tests to follow Signed-off-by: d10n <[email protected]> * Reuse existing merge_dicts_deep Signed-off-by: d10n <[email protected]> * Fix SyntaxWarning: invalid escape sequence '\[' Signed-off-by: d10n <[email protected]> * Fix lint errors Signed-off-by: d10n <[email protected]> --------- Signed-off-by: d10n <[email protected]> Co-authored-by: Jennifer Power <[email protected]>
oscal-compass · Nov 19, 2024 · 9f1edcf · 9f1edcf
1 parent 65d7dce
commit 9f1edcf
Show file tree

Hide file tree

Showing 5 changed files with 55 additions and 11 deletions.
diff --git a/tests/trestle/core/control_io_test.py b/tests/trestle/core/control_io_test.py
@@ -30,7 +30,7 @@
 from trestle.common.model_utils import ModelUtils
 from trestle.core.catalog.catalog_interface import CatalogInterface
 from trestle.core.control_context import ContextPurpose, ControlContext
-from trestle.core.control_interface import ControlInterface, ParameterRep
+from trestle.core.control_interface import ComponentImpInfo, ControlInterface, ParameterRep
 from trestle.core.control_reader import ControlReader
 from trestle.core.control_writer import ControlWriter
 from trestle.core.markdown.control_markdown_node import ControlMarkdownNode, tree_context
@@ -45,7 +45,7 @@
 case_3 = 'indent end abrupt'
 case_4 = 'no items'
 
-control_text = """---
+control_text = r"""---
 x-trestle-global:
   sort-id: xy-09
 ---
@@ -131,7 +131,7 @@ def test_read_write_controls(
     part_b3 = common.Part(id='ac-1_smt.b.3', name='item', prose='b.3 prose', props=[prop])
     prop.value = 'c'
     part_c = common.Part(id='ac-1_smt.c', name='item', prose='c prose', props=[prop])
-    sec_1_text = """
+    sec_1_text = r"""
 General comment
 on separate lines
 
@@ -395,7 +395,47 @@ def test_write_control_header_params(overwrite_header_values, tmp_path: pathlib.
     assert test_utils.controls_equivalent(orig_control_read, new_control_read)
 
 
-statement_text = """
+def test_merge_control_update(tmp_path: pathlib.Path, testdata_dir: pathlib.Path) -> None:
+    """Test merging of control header params after spec update."""
+    src_control_path = pathlib.Path(testdata_dir / 'author/controls/control_with_components.md')
+    control_path = tmp_path / 'ac-1.md'
+    shutil.copyfile(src_control_path, control_path)
+    orig_control_read, group_title = ControlReader.read_control(control_path, False)
+    assert group_title == 'Access Control'
+    context = ControlContext.generate(ContextPurpose.COMPONENT, True, tmp_path, tmp_path, True)
+    # Given updated template comp_dict from the component definition json
+    context.comp_dict = {
+        'This System': {
+            '': ComponentImpInfo(
+                prose='', rules=[], props=[], status=common.ImplementationStatus(state='planned', remarks=None)
+            ),
+            'a.': ComponentImpInfo(
+                prose='Text for fancy thing component',
+                rules=[],
+                props=[],
+                status=common.ImplementationStatus(state='planned', remarks=None)
+            ),
+            'c.': ComponentImpInfo(
+                prose='Just for the default component',
+                rules=[],
+                props=[],
+                status=common.ImplementationStatus(state='planned', remarks=None)
+            ),
+            'd.': ComponentImpInfo(
+                prose='Example extra component',
+                rules=[],
+                props=[],
+                status=common.ImplementationStatus(state='planned', remarks=None)
+            )
+        }
+    }
+    control_writer = ControlWriter()
+    control_writer.write_control_for_editing(context, orig_control_read, tmp_path, group_title, {}, [])
+    assert context.comp_dict['This System']['c.'].status.state == 'operational', 'State must be merged'
+    assert context.comp_dict['This System']['d.'].status.state == 'planned', 'New template state must be merged'
+
+
+statement_text = r"""
 
 
 # xy-9 - \[My Group Title\] Fancy Control

diff --git a/trestle/core/catalog/catalog_interface.py b/trestle/core/catalog/catalog_interface.py
@@ -904,5 +904,5 @@ def generate_control_rule_info(self, part_id_map: Dict[str, Dict[str, str]], con
         if len(dup_comp_uuids) > 0:
             # throw an exception if there are repeated component uuids
             for comp_uuid in dup_comp_uuids:
-                logger.error(f'Component uuid { comp_uuid } is duplicated')
+                logger.error(f'Component uuid {comp_uuid} is duplicated')
             raise TrestleError('Component uuids cannot be duplicated between different component definitions')
diff --git a/trestle/core/control_interface.py b/trestle/core/control_interface.py
@@ -109,8 +109,8 @@ class ControlInterface:
 
     @staticmethod
     def _wrap_label(label: str):
-        l_side = '\['
-        r_side = '\]'
+        l_side = r'\['
+        r_side = r'\]'
         wrapped = '' if label == '' else f'{l_side}{label}{r_side}'
         return wrapped
 
@@ -591,6 +591,8 @@ def merge_dicts_deep(
         New items are always added from src to dest.
         Items present in both will be overriden dest if overwrite_header_values is True.
         """
+        if src is None:
+            return
         for key in src.keys():
             if key in dest:
                 if depth and level == depth:

diff --git a/trestle/core/control_writer.py b/trestle/core/control_writer.py
@@ -66,7 +66,7 @@ def _add_control_statement(self, control: cat.Control, group_title: str, print_g
         control_title = control.title
 
         if print_group_title:
-            group_name = ' \[' + group_title + '\]'
+            group_name = r' \[' + group_title + r'\]'
 
         title = f'{control_id} -{group_name} {control_title}'
 
@@ -516,8 +516,10 @@ def write_control_for_editing(
         control_file = dest_path / (control.id + const.MARKDOWN_FILE_EXT)
         # read the existing markdown header and content if it exists
         md_header, comp_dict = ControlReader.read_control_info_from_md(control_file, context)
-        # replace the memory comp_dict with the md one if control exists
+        # Merge the memory comp_dict with the md one if control exists
         if comp_dict:
+            template_comp_dict = context.comp_dict
+            ControlInterface.merge_dicts_deep(comp_dict, template_comp_dict, False)
             context.comp_dict = comp_dict
 
         header_comment_dict = {const.TRESTLE_ADD_PROPS_TAG: const.YAML_PROPS_COMMENT}

diff --git a/trestle/transforms/implementations/tanium.py b/trestle/transforms/implementations/tanium.py
@@ -110,7 +110,7 @@ def transform(self, blob: str) -> Results:
         results.__root__ = tanium_oscal_factory.results
         ts1 = datetime.datetime.now()
         self._analysis = tanium_oscal_factory.analysis
-        self._analysis.append(f'transform time: {ts1-ts0}')
+        self._analysis.append(f'transform time: {ts1 - ts0}')
         return results
 
 
@@ -455,7 +455,7 @@ def _batch_observations(self, index: int) -> Dict[str, List[Observation]]:
         start = index * batch_size
         end = (index + 1) * batch_size
         end = min(end, len(self._rule_use_list))
-        logger.debug(f'start: {start} end: {end-1}')
+        logger.debug(f'start: {start} end: {end - 1}')
         # process just the one chunk
         for i in range(start, end):
             rule_use = self._rule_use_list[i]