PSCE-258 - feat: adds filtering by profile to AuthoredComponentDefini…

…tion create… (#62)

* feat: adds filtering by profile to AuthoredComponentDefinition create_new_default

Signed-off-by: Jennifer Power <[email protected]>

* test: fixes flaky test under test_rule_transform_task.py

Signed-off-by: Jennifer Power <[email protected]>

---------

Signed-off-by: Jennifer Power <[email protected]>

tests/trestlebot/tasks/authored/test_compdef.py

@@ -20,6 +20,10 @@
 import re
 
 import pytest
+from trestle.common.model_utils import ModelUtils
+from trestle.core.catalog.catalog_interface import CatalogInterface
+from trestle.core.profile_resolver import ProfileResolver
+from trestle.oscal import profile as prof
 
 from tests import testutils
 from trestlebot.const import RULES_VIEW_DIR, YAML_EXTENSION
@@ -83,6 +87,43 @@ def test_create_new_default(tmp_trestle_dir: str) -> None:
     assert rule.profile.include_controls[0].id == "ac-5"
 
 
+def test_create_new_default_with_filter(tmp_trestle_dir: str) -> None:
+    """Test creating new default component definition with filter"""
+    # Prepare the workspace
+    trestle_root = pathlib.Path(tmp_trestle_dir)
+    _ = testutils.setup_for_profile(trestle_root, test_prof, "")
+    authored_comp = AuthoredComponentDefinition(tmp_trestle_dir)
+
+    profile_path = ModelUtils.get_model_path_for_name_and_class(
+        trestle_root, test_prof, prof.Profile
+    )
+
+    catalog = ProfileResolver.get_resolved_profile_catalog(
+        trestle_root, profile_path=profile_path
+    )
+
+    catalog_interface = CatalogInterface(catalog)
+    catalog_interface.delete_control("ac-5")
+
+    authored_comp.create_new_default(
+        test_prof, test_comp, "test", "My desc", "service", catalog_interface
+    )
+
+    rules_view_dir = trestle_root / RULES_VIEW_DIR
+    assert rules_view_dir.exists()
+
+    compdef_dir = rules_view_dir / test_comp
+    assert compdef_dir.exists()
+
+    comp_dir = compdef_dir / "test"
+    assert comp_dir.exists()
+
+    # Verity that the number of rules YAML files has been reduced
+    # from 12 to 11.
+    yaml_files = list(comp_dir.glob(f"*{YAML_EXTENSION}"))
+    assert len(yaml_files) == 11
+
+
 def test_create_new_default_no_profile(tmp_trestle_dir: str) -> None:
     """Test creating new default component definition successfully"""
     # Prepare the workspace

tests/trestlebot/tasks/test_rule_transform_task.py

@@ -56,20 +56,24 @@ def test_rule_transform_task(tmp_trestle_dir: str) -> None:
     assert orig_comp.components is not None
     assert len(orig_comp.components) == 2
 
-    component = orig_comp.components[0]
+    component = next(
+        (comp for comp in orig_comp.components if comp.title == "Component 2"), None
+    )
 
+    assert component is not None
     assert component.props is not None
-    assert component.title == "Component 2"
     assert len(component.props) == 2
     assert component.props[0].name == RULE_ID
     assert component.props[0].value == "example_rule_2"
     assert component.props[1].name == RULE_DESCRIPTION
     assert component.props[1].value == "My rule description for example rule 2"
 
-    component = orig_comp.components[1]
+    component = next(
+        (comp for comp in orig_comp.components if comp.title == "Component 1"), None
+    )
 
+    assert component is not None
     assert component.props is not None
-    assert component.title == "Component 1"
     assert len(component.props) == 5
     assert component.props[0].name == RULE_ID
     assert component.props[0].value == "example_rule_1"

trestlebot/tasks/authored/compdef.py

@@ -18,7 +18,7 @@
 
 import os
 import pathlib
-from typing import List
+from typing import Callable, List, Optional
 
 import trestle.common.const as const
 import trestle.oscal.profile as prof
@@ -103,6 +103,7 @@ def create_new_default(
         comp_title: str,
         comp_description: str,
         comp_type: str,
+        filter_controls: Optional[CatalogInterface] = None,
     ) -> None:
         """
         Create the new component definition with default info.
@@ -113,6 +114,7 @@ def create_new_default(
             comp_title: Title of the component
             comp_description: Description of the component
             comp_type: Type of the component
+            filter_controls: Optional catalog to filter the controls to include from the profile
 
         Notes:
             The beginning of the Component Definition workflow is to create a new
@@ -137,10 +139,29 @@ def create_new_default(
         )
 
         rules_view_builder = RulesViewBuilder(trestle_root)
-        rules_view_builder.add_rules_for_profile(existing_profile_path, component_info)
+
+        filter_func: Optional[Callable[[str], bool]] = None
+        if filter_controls is not None:
+            filter_func = FilterByCatalog(filter_controls)
+
+        rules_view_builder.add_rules_for_profile(
+            existing_profile_path, component_info, filter_func
+        )
         rules_view_builder.write_to_yaml(rule_dir)
 
 
+class FilterByCatalog:
+    """Filter controls by catalog."""
+
+    def __init__(self, catalog: CatalogInterface) -> None:
+        """Initialize."""
+        self._catalog = catalog
+
+    def __call__(self, control_id: str) -> bool:
+        """Filter controls by catalog."""
+        return control_id in self._catalog.get_control_ids()
+
+
 class RulesViewBuilder:
     """Write TrestleRule objects to YAML files in rules view."""
 
@@ -151,16 +172,29 @@ def __init__(self, trestle_root: pathlib.Path) -> None:
         self._yaml_transformer = FromRulesYAMLTransformer()
 
     def add_rules_for_profile(
-        self, profile_path: pathlib.Path, component_info: ComponentInfo
+        self,
+        profile_path: pathlib.Path,
+        component_info: ComponentInfo,
+        criteria: Optional[Callable[[str], bool]] = None,
     ) -> None:
-        """Add rules for a profile to the builder."""
+        """
+        Add rules for a profile to the builder.
+
+        Args:
+            profile_path: Path to the profile
+            component_info: Component info to use for the rules
+            criteria: Optional criteria to filter the controls to include in the rules
+        """
         catalog = ProfileResolver.get_resolved_profile_catalog(
             self._trestle_root, profile_path=profile_path
         )
 
         controls = CatalogInterface.get_control_ids_from_catalog(catalog)
 
         for control_id in controls:
+            if criteria is not None and not criteria(control_id):
+                continue
+
             rule = TrestleRule(
                 component=component_info,
                 name=f"rule-{control_id}",