Updates to tests and specification.

Author: rlhagerm

python/example_code/controltower/controltower_wrapper.py

@@ -303,11 +303,10 @@ def list_landing_zones(self):
     # snippet-end:[python.example_code.controltower.ListLandingZones]
 
     # snippet-start:[python.example_code.controltower.ListEnabledBaselines]
-    def list_enabled_baselines(self, target_identifier):
+    def list_enabled_baselines(self):
         """
-        Lists all enabled baselines for a specific target.
+        Lists all enabled baselines.
 
-        :param target_identifier: The identifier of the target (e.g., OU ARN).
         :return: List of enabled baselines.
         :raises ClientError: If the listing operation fails.
         """
@@ -387,15 +386,16 @@ def disable_baseline(self, enabled_baseline_identifier):
 
             return response['operationIdentifier']
         except ClientError as err:
-            if err.response["Error"]["Code"] == "ResourceNotFoundException":
-                logger.error("Target not found.")
+            if err.response["Error"]["Code"] == "ConflictException":
+                print(f"Conflict disabling baseline: {err.response['Error']['Message']}. Skipping disable step." )
+                return None
             else:
                 logger.error(
                     "Couldn't disable baseline. Here's why: %s: %s",
                     err.response["Error"]["Code"],
                     err.response["Error"]["Message"]
                 )
-            raise
+                raise
     # snippet-end:[python.example_code.controltower.DisableBaseline]
 
     # snippet-start:[python.example_code.controltower.ListEnabledControls]

python/example_code/controltower/scenario_controltower.py

@@ -82,6 +82,7 @@ def run_scenario(self):
                 self.ou_id = sandbox_ou_id
 
         # List and Enable Baseline.
+        print("\nManaging Baselines:")
         control_tower_baseline = None
         identity_center_baseline = None
         baselines = self.controltower_wrapper.list_baselines()
@@ -93,9 +94,7 @@ def run_scenario(self):
 
         if self.use_landing_zone:
             print("\nListing enabled baselines:")
-            enabled_baselines = self.controltower_wrapper.list_enabled_baselines(
-                self.ou_arn
-            )
+            enabled_baselines = self.controltower_wrapper.list_enabled_baselines()
             for baseline in enabled_baselines:
                 # If the Identity Center baseline is enabled, the identifier must be used for other baselines.
                 if 'baseline/LN25R72TTG6IGPTQ' in baseline['baselineIdentifier']:
@@ -107,25 +106,27 @@ def run_scenario(self):
                     q.is_yesno,
             ):
                 print("\nEnabling Control Tower Baseline.")
+                ic_baseline_arn = identity_center_baseline['arn'] if identity_center_baseline else None
                 baseline_arn = self.controltower_wrapper.enable_baseline(
                     self.ou_arn,
-                    identity_center_baseline['arn'],
+                    ic_baseline_arn,
                     control_tower_baseline['arn'],
                     '4.0'
                 )
                 if baseline_arn:
                     print(f"Enabled baseline ARN: {baseline_arn}")
                 else:
+                    # Find the enabled baseline so we can reset it.
                     for enabled_baseline in enabled_baselines:
-                        if enabled_baseline['arn'] == control_tower_baseline['arn']:
-                            control_tower_baseline = baseline
+                        if enabled_baseline['baselineIdentifier'] == control_tower_baseline['arn']:
+                            baseline_arn = enabled_baseline['arn']
                     print("No change, the selected baseline was already enabled.")
 
                 if q.ask(
                         f"Do you want to reset the Control Tower Baseline? (y/n) ",
                         q.is_yesno,
                 ):
-                    print("\nResetting Control Tower Baseline.")
+                    print(f"\nResetting Control Tower Baseline. {baseline_arn}")
                     operation_id = self.controltower_wrapper.reset_enabled_baseline(
                         baseline_arn
                     )
@@ -142,23 +143,22 @@ def run_scenario(self):
                     print(f"\nDisabled baseline operation id {operation_id}.")
 
         # List and Enable Controls.
-        print("Managing Controls:")
+        print("\nManaging Controls:")
         controls = self.controltower_wrapper.list_controls()
         print("\nListing first 5 available Controls:")
         for i, control in enumerate(controls[:5], 1):
-            print(f"{i}. {control['Name']}")
+            print(f"{i}. {control['Name']} - {control['Arn']}")
 
         if self.use_landing_zone:
-
-            enabled_controls = self.controltower_wrapper.list_enabled_controls()
+            target_ou = self.ou_arn
+            enabled_controls = self.controltower_wrapper.list_enabled_controls(target_ou)
             print("\nListing enabled controls:")
-            for i, control in enabled_controls:
-                print(f"{i}. {control['Name']}")
+            for i, control in enumerate(enabled_controls, 1):
+                print(f"{i}. {control['controlIdentifier']}")
 
             # Enable first non-enabled control as an example.
-            enabled_control_arns = [control['Arn'] for control in enabled_controls]
+            enabled_control_arns = [control['arn'] for control in enabled_controls]
             control_arn = next(control['Arn'] for control in controls if control['Arn'] not in enabled_control_arns)
-            target_ou = self.ou_arn
 
             if control_arn and q.ask(
                 f"Do you want to enable the control {control_arn}? (y/n) ",
@@ -182,10 +182,10 @@ def run_scenario(self):
                     control_arn, target_ou)
                 print(f"Disable operation ID: {operation_id}")
 
-            print("This concludes the control tower scenario.")
+        print("\nThis concludes the example scenario.")
 
-            print("Thanks for watching!")
-            print("-" * 88)
+        print("Thanks for watching!")
+        print("-" * 88)
 
     def setup_organization(self):
         """

python/example_code/controltower/test/test_scenario_disable_baseline.py

@@ -0,0 +1,62 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Unit tests for the disable_baseline method in controltower_wrapper.py.
+"""
+
+import pytest
+from botocore.exceptions import ClientError
+
+class MockManager:
+    def __init__(self, stub_runner, scenario_data, input_mocker):
+        self.scenario_data = scenario_data
+        self.ou_arn = "arn:aws:organizations::123456789012:ou/o-exampleorgid/ou-exampleouid"
+        self.baseline_arn = "arn:aws:controltower:us-east-1:123456789012:baseline/AWSControlTowerBaseline/enabled"
+        self.operation_id = "op-1234567890abcdef0"
+        
+        self.stub_runner = stub_runner
+
+    def setup_stubs(self, error, stop_on, controltower_stubber):
+        with self.stub_runner(error, stop_on) as runner:
+            runner.add(
+                controltower_stubber.stub_disable_baseline,
+                self.ou_arn,
+                self.baseline_arn,
+                self.operation_id
+            )
+
+
+@pytest.fixture
+def mock_mgr(stub_runner, scenario_data, input_mocker):
+    return MockManager(stub_runner, scenario_data, input_mocker)
+
+@pytest.mark.integ
+def test_disable_baseline(mock_mgr, capsys):
+    mock_mgr.setup_stubs(None, None, mock_mgr.scenario_data.controltower_stubber)
+    
+    # Test disabling a baseline
+    operation_id = mock_mgr.scenario_data.scenario.controltower_wrapper.disable_baseline(
+        mock_mgr.ou_arn,
+        mock_mgr.baseline_arn
+    )
+    
+    # Verify the results
+    assert operation_id == mock_mgr.operation_id
+
+
+@pytest.mark.parametrize(
+    "error, stop_on_index",
+    [
+        ("TESTERROR-stub_disable_baseline", 0),
+    ],
+)
+@pytest.mark.integ
+def test_disable_baseline_error(mock_mgr, caplog, error, stop_on_index):
+    mock_mgr.setup_stubs(error, stop_on_index, mock_mgr.scenario_data.controltower_stubber)
+    
+    with pytest.raises(ClientError) as exc_info:
+        mock_mgr.scenario_data.scenario.controltower_wrapper.disable_baseline(
+            mock_mgr.ou_arn,
+            mock_mgr.baseline_arn
+        )

python/example_code/controltower/test/test_scenario_enabled_baselines.py

@@ -0,0 +1,73 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Unit tests for the list_enabled_baselines method in controltower_wrapper.py.
+"""
+
+import pytest
+from botocore.exceptions import ClientError
+
+class MockManager:
+    def __init__(self, stub_runner, scenario_data, input_mocker):
+        self.scenario_data = scenario_data
+        self.ou_arn = "arn:aws:organizations::123456789012:ou/o-exampleorgid/ou-exampleouid"
+        
+        self.enabled_baselines = [
+            {
+                "baselineArn": "arn:aws:controltower:us-east-1:123456789012:baseline/AWSControlTowerBaseline/enabled",
+                "baselineVersion": "4.0",
+                "baselineName": "AWSControlTowerBaseline"
+            },
+            {
+                "baselineArn": "arn:aws:controltower:us-east-1:123456789012:baseline/OtherBaseline/enabled",
+                "baselineVersion": "2.0",
+                "baselineName": "OtherBaseline"
+            }
+        ]
+        
+        self.stub_runner = stub_runner
+
+    def setup_stubs(self, error, stop_on, controltower_stubber):
+        with self.stub_runner(error, stop_on) as runner:
+            runner.add(
+                controltower_stubber.stub_list_enabled_baselines,
+                self.ou_arn,
+                self.enabled_baselines
+            )
+
+
+@pytest.fixture
+def mock_mgr(stub_runner, scenario_data, input_mocker):
+    return MockManager(stub_runner, scenario_data, input_mocker)
+
+@pytest.mark.integ
+def test_list_enabled_baselines(mock_mgr, capsys):
+    mock_mgr.setup_stubs(None, None, mock_mgr.scenario_data.controltower_stubber)
+    
+    # Test listing enabled baselines
+    enabled_baselines = mock_mgr.scenario_data.scenario.controltower_wrapper.list_enabled_baselines(
+        mock_mgr.ou_arn
+    )
+    
+    # Verify the results
+    assert len(enabled_baselines) == 2
+    assert enabled_baselines[0]["baselineName"] == "AWSControlTowerBaseline"
+    assert enabled_baselines[0]["baselineVersion"] == "4.0"
+    assert enabled_baselines[1]["baselineName"] == "OtherBaseline"
+
+
+@pytest.mark.parametrize(
+    "error, stop_on_index",
+    [
+        ("TESTERROR-stub_list_enabled_baselines", 0),
+    ],
+)
+@pytest.mark.integ
+def test_list_enabled_baselines_error(mock_mgr, caplog, error, stop_on_index):
+    mock_mgr.setup_stubs(error, stop_on_index, mock_mgr.scenario_data.controltower_stubber)
+    
+    with pytest.raises(ClientError) as exc_info:
+        mock_mgr.scenario_data.scenario.controltower_wrapper.list_enabled_baselines(
+            mock_mgr.ou_arn
+        )

python/example_code/controltower/test/test_scenario_enabled_controls.py

@@ -0,0 +1,73 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Unit tests for the list_enabled_controls method in controltower_wrapper.py.
+"""
+
+import pytest
+from botocore.exceptions import ClientError
+
+class MockManager:
+    def __init__(self, stub_runner, scenario_data, input_mocker):
+        self.scenario_data = scenario_data
+        self.ou_arn = "arn:aws:organizations::123456789012:ou/o-exampleorgid/ou-exampleouid"
+        
+        self.enabled_controls = [
+            {
+                "controlIdentifier": "arn:aws:controlcatalog:us-east-1:123456789012:control/aws-control-1234",
+                "controlName": "TestControl1",
+                "controlStatus": "ENABLED"
+            },
+            {
+                "controlIdentifier": "arn:aws:controlcatalog:us-east-1:123456789012:control/aws-control-5678",
+                "controlName": "TestControl2",
+                "controlStatus": "ENABLED"
+            }
+        ]
+        
+        self.stub_runner = stub_runner
+
+    def setup_stubs(self, error, stop_on, controltower_stubber):
+        with self.stub_runner(error, stop_on) as runner:
+            runner.add(
+                controltower_stubber.stub_list_enabled_controls,
+                self.ou_arn,
+                self.enabled_controls
+            )
+
+
+@pytest.fixture
+def mock_mgr(stub_runner, scenario_data, input_mocker):
+    return MockManager(stub_runner, scenario_data, input_mocker)
+
+@pytest.mark.integ
+def test_list_enabled_controls(mock_mgr, capsys):
+    mock_mgr.setup_stubs(None, None, mock_mgr.scenario_data.controltower_stubber)
+    
+    # Test listing enabled controls
+    enabled_controls = mock_mgr.scenario_data.scenario.controltower_wrapper.list_enabled_controls(
+        mock_mgr.ou_arn
+    )
+    
+    # Verify the results
+    assert len(enabled_controls) == 2
+    assert enabled_controls[0]["controlName"] == "TestControl1"
+    assert enabled_controls[0]["controlStatus"] == "ENABLED"
+    assert enabled_controls[1]["controlName"] == "TestControl2"
+
+
+@pytest.mark.parametrize(
+    "error, stop_on_index",
+    [
+        ("TESTERROR-stub_list_enabled_controls", 0),
+    ],
+)
+@pytest.mark.integ
+def test_list_enabled_controls_error(mock_mgr, caplog, error, stop_on_index):
+    mock_mgr.setup_stubs(error, stop_on_index, mock_mgr.scenario_data.controltower_stubber)
+    
+    with pytest.raises(ClientError) as exc_info:
+        mock_mgr.scenario_data.scenario.controltower_wrapper.list_enabled_controls(
+            mock_mgr.ou_arn
+        )