Throw specific Exception instead of generic one

This to solve check flake8 B017 Also disable B028 flake8-bugbear lint check recently implemented. The check is going to be disabled by the flake8-bugbear developers. See Rename B028 to B907, making it optional/opinionated. PyCQA/flake8-bugbear#333 Fix color formatting on flake8 style offenses. Signed-off-by: Luca Carrogu <carrogu@amazon.com>
lukeseawalker · Jan 17, 2023 · 77b545a · 77b545a
1 parent 554c96f
commit 77b545a
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 25 deletions.
diff --git a/.flake8 b/.flake8
@@ -11,7 +11,10 @@ ignore =
     # W503: line break before binary operator => Conflicts with black style.
     W503,
     # D413: Missing blank line after last section
-    D413
+    D413,
+    # B028: Consider replacing f"'{foo}'" with f"{foo!r}".
+    # Currently being disabled by flake8-bugbear. See https://github.com/PyCQA/flake8-bugbear/pull/333
+    B028
 exclude =
     .tox,
     .git,
@@ -28,4 +31,4 @@ max-complexity = 10
 max-line-length = 120
 import-order-style = google
 application-import-names = flake8
-format = ${cyan}%(path)s${reset}:${yellow_bold}%(row)d${reset}:${green_bold}%(col)d${reset}: ${red_bold}%(code)s${reset} %(text)s
+format = %(cyan)s%(path)s%(reset)s:%(bold)s%(yellow)s%(row)d%(reset)s:%(bold)s%(green)s%(col)d%(reset)s: %(bold)s%(red)s%(code)s%(reset)s %(text)s
diff --git a/cookbooks/aws-parallelcluster-install/files/default/clusterstatusmgtd/clusterstatusmgtd.py b/cookbooks/aws-parallelcluster-install/files/default/clusterstatusmgtd/clusterstatusmgtd.py
@@ -206,6 +206,11 @@ class ConditionalStatusUpdateFailedError(Exception):
 
         pass
 
+    class FleetDataNotFoundError(Exception):
+        """Raised when compute fleet data cannot be found in db table."""
+
+        pass
+
     def __init__(self, table_name, boto3_config, region):
         self._ddb_resource = boto3.resource("dynamodb", region_name=region, config=boto3_config)
         self._table = self._ddb_resource.Table(table_name)
@@ -216,7 +221,7 @@ def get_status(self):  # noqa: D102
             Key={"Id": self.DB_KEY},
         )
         if not compute_fleet_item or "Item" not in compute_fleet_item:
-            raise Exception("COMPUTE_FLEET data not found in db table")
+            raise ComputeFleetStatusManager.FleetDataNotFoundError("COMPUTE_FLEET data not found in db table")
 
         log.debug("Found COMPUTE_FLEET data (%s)", compute_fleet_item)
         return compute_fleet_item["Item"].get(self.DB_DATA)
@@ -323,6 +328,11 @@ def __init__(self, config):
         self._compute_fleet_data = {}
         self.set_config(config)
 
+    class ClusterStatusUpdateEventError(Exception):
+        """Raised when there is a failure in updating the status due to an error on update event handler execution."""
+
+        pass
+
     def set_config(self, config):  # noqa: D102
         if self._config != config:
             log.info("Applying new clusterstatusmgtd config: %s", config)
@@ -366,7 +376,7 @@ def _call_update_event(self):
             _write_json_to_file(self._config.computefleet_status_path, compute_fleet_data)
         except Exception as e:
             log.error("Update event handler failed during fleet status translation: %s", e)
-            raise
+            raise ClusterStatusManager.ClusterStatusUpdateEventError(e)
 
         cinc_log_file = "/var/log/chef-client.log"
         log.info("Calling update event handler, log can be found at %s", cinc_log_file)
@@ -384,9 +394,9 @@ def _call_update_event(self):
         )
         try:
             _run_command(cmd, self._config.update_event_timeout_minutes)
-        except Exception:
+        except Exception as e:
             log.error("Update event handler failed. Check log file %s", cinc_log_file)
-            raise
+            raise ClusterStatusManager.ClusterStatusUpdateEventError(e)
 
     def _update_status(self, request_status, in_progress_status, final_status):
         if self._compute_fleet_status == request_status:

diff --git a/...allelcluster-slurm/files/default/head_node_slurm/slurm/pcluster_fleet_config_generator.py b/...allelcluster-slurm/files/default/head_node_slurm/slurm/pcluster_fleet_config_generator.py
@@ -25,6 +25,12 @@ class CriticalError(Exception):
     pass
 
 
+class ConfigurationFieldNotFoundError(Exception):
+    """Field not found in configuration."""
+
+    pass
+
+
 def generate_fleet_config_file(output_file, input_file):
     """
     Generate configuration file used by Fleet Manager in node daemon package.
@@ -87,7 +93,7 @@ def generate_fleet_config_file(output_file, input_file):
                     }
 
                 else:
-                    raise Exception(
+                    raise ConfigurationFieldNotFoundError(
                         "Instances or InstanceType field not found "
                         f"in queue: {queue}, compute resource: {compute_resource} configuration"
                     )

diff --git a/test/unit/clusterstatusmgtd/test_clusterstatusmgtd.py b/test/unit/clusterstatusmgtd/test_clusterstatusmgtd.py
@@ -123,21 +123,18 @@ def compute_fleet_status_manager(self, mocker):
         return status_manager
 
     @pytest.mark.parametrize(
-        "get_item_response, expected_status",
+        "get_item_response, expected_exception, expected_status",
         [
-            ({"Item": {"Id": "COMPUTE_FLEET", "Data": {"status": "RUNNING"}}}, {"status": "RUNNING"}),
-            (
-                Exception,
-                None,
-            ),
+            ({"Item": {"Id": "COMPUTE_FLEET", "Data": {"status": "RUNNING"}}}, None, {"status": "RUNNING"}),
+            ({"NoData": "NoValue"}, ComputeFleetStatusManager.FleetDataNotFoundError, Exception()),
         ],
         ids=["success", "exception"],
     )
-    def test_get_status(self, compute_fleet_status_manager, get_item_response, expected_status):
+    def test_get_status(self, compute_fleet_status_manager, get_item_response, expected_exception, expected_status):
         """Test get_status method."""
-        if get_item_response is Exception:
+        if isinstance(expected_status, Exception):
             compute_fleet_status_manager._table.get_item.side_effect = get_item_response
-            with pytest.raises(Exception):
+            with pytest.raises(expected_exception):
                 compute_fleet_status_manager.get_status()
         else:
             compute_fleet_status_manager._table.get_item.return_value = get_item_response
@@ -398,7 +395,7 @@ def test_call_update_event(self, mocker, status, translated_status, exception):
         run_command_mock = mocker.patch("clusterstatusmgtd._run_command")
         if isinstance(exception, Exception):
             run_command_mock.side_effect = exception
-            with pytest.raises(Exception):
+            with pytest.raises(ClusterStatusManager.ClusterStatusUpdateEventError):
                 clusterstatus_manager._call_update_event()
         else:
             file_writer_mock = mocker.mock_open()

diff --git a/test/unit/slurm/test_fleet_config_generator.py b/test/unit/slurm/test_fleet_config_generator.py
@@ -12,34 +12,42 @@
 
 import pytest
 from assertpy import assert_that
-from slurm.pcluster_fleet_config_generator import generate_fleet_config_file
+from slurm.pcluster_fleet_config_generator import (
+    ConfigurationFieldNotFoundError,
+    CriticalError,
+    generate_fleet_config_file,
+)
 
 
 @pytest.mark.parametrize(
-    "cluster_config, expected_message",
+    "cluster_config, expected_exception, expected_message",
     [
-        ({}, "Unable to find key 'Scheduling' in the configuration file"),
-        ({"Scheduling": {}}, "Unable to find key 'SlurmQueues' in the configuration file"),
-        ({"Scheduling": {"SlurmQueues": []}}, None),
+        ({}, CriticalError, "Unable to find key 'Scheduling' in the configuration file"),
+        ({"Scheduling": {}}, CriticalError, "Unable to find key 'SlurmQueues' in the configuration file"),
+        ({"Scheduling": {"SlurmQueues": []}}, None, None),
         (
             {"Scheduling": {"SlurmQueues": [{"ComputeResources": []}]}},
+            CriticalError,
             "Unable to find key 'Name' in the configuration file",
         ),
         (
             {"Scheduling": {"SlurmQueues": [{"Name": "q1"}]}},
+            CriticalError,
             "Unable to find key 'CapacityType' in the configuration of queue: q1",
         ),
         (
             {"Scheduling": {"SlurmQueues": [{"Name": "q1", "CapacityType": "ONDEMAND"}]}},
+            CriticalError,
             "Unable to find key 'ComputeResources' in the configuration of queue: q1",
         ),
-        ({"Scheduling": {"SlurmQueues": [{"Name": "q1", "CapacityType": "SPOT", "ComputeResources": []}]}}, None),
+        ({"Scheduling": {"SlurmQueues": [{"Name": "q1", "CapacityType": "SPOT", "ComputeResources": []}]}}, None, None),
         (
             {
                 "Scheduling": {
                     "SlurmQueues": [{"Name": "q1", "CapacityType": "SPOT", "ComputeResources": [{"Instances": []}]}]
                 }
             },
+            CriticalError,
             "Unable to find key 'Name' in the configuration of queue: q1",
         ),
         (
@@ -55,6 +63,7 @@
                     ]
                 }
             },
+            ConfigurationFieldNotFoundError,
             "Instances or InstanceType field not found in queue: q1, compute resource: cr1 configuration",
         ),
         (
@@ -74,6 +83,7 @@
                 }
             },
             None,
+            None,
         ),
         (
             {
@@ -92,6 +102,7 @@
                 }
             },
             None,
+            None,
         ),
         (
             {
@@ -114,6 +125,7 @@
                 }
             },
             None,
+            None,
         ),
         (
             {
@@ -128,6 +140,7 @@
                     ]
                 }
             },
+            CriticalError,
             "Unable to find key 'SpotPrice' in the configuration of queue: q1, compute resource: cr1",
         ),
         (
@@ -146,6 +159,7 @@
                 }
             },
             None,
+            None,
         ),
         (
             {
@@ -161,6 +175,7 @@
                     ]
                 }
             },
+            CriticalError,
             "Unable to find key 'Networking' in the configuration of queue: q1, compute resource: cr1",
         ),
         (
@@ -178,16 +193,17 @@
                     ]
                 }
             },
+            CriticalError,
             "Unable to find key 'SubnetIds' in the configuration of queue: q1, compute resource: cr1",
         ),
     ],
 )
-def test_generate_fleet_config_file_error_cases(mocker, tmpdir, cluster_config, expected_message):
+def test_generate_fleet_config_file_error_cases(mocker, tmpdir, cluster_config, expected_exception, expected_message):
     mocker.patch("slurm.pcluster_fleet_config_generator._load_cluster_config", return_value=cluster_config)
     output_file = f"{tmpdir}/fleet-config.json"
 
     if expected_message:
-        with pytest.raises(Exception, match=expected_message):
+        with pytest.raises(expected_exception, match=expected_message):
             generate_fleet_config_file(output_file, input_file="fake")
     else:
         generate_fleet_config_file(output_file, input_file="fake")