Unify CLI config-failure exit codes via a single ConfigError catch

main.py

@@ -162,7 +162,7 @@ def handle_aws(args):
         config = load_config(args.config)
         if not config:
             console.print("[red]Invalid or missing AWS configuration file.[/red]")
-            return
+            raise ConfigError
 
         # Handle name field logic (priority: --name > config name > fallback)
         if args.name:
@@ -182,25 +182,16 @@ def handle_aws(args):
         assessment_type = require_env_int(
             "ESC_ASSESSMENT_TYPE", "assessment type (1 or 2)", {1, 2}
         )
-        try:
-            if args.profile:
-                provider_details = _aws_provider_from_profile(args.profile)
-            else:
-                provider_details = _aws_provider_from_env()
-        except ConfigError:
-            sys.exit(codes.CONFIG)
-    elif args.profile:
-        try:
+        if args.profile:
             provider_details = _aws_provider_from_profile(args.profile)
-        except ConfigError:
-            return
+        else:
+            provider_details = _aws_provider_from_env()
+    elif args.profile:
+        provider_details = _aws_provider_from_profile(args.profile)
         exit_strategy, assessment_type = prompt_required_inputs()
     else:
         exit_strategy, assessment_type = prompt_required_inputs()
-        try:
-            provider_details = _aws_provider_from_prompt()
-        except ConfigError:
-            return
+        provider_details = _aws_provider_from_prompt()
 
     config = build_config(
         cloud_provider, exit_strategy, assessment_type, provider_details, args
@@ -364,7 +355,7 @@ def handle_azure(args):
         config = load_config(args.config)
         if not config:
             console.print("[red]Invalid or missing Azure configuration file.[/red]")
-            return
+            raise ConfigError
 
         # Handle name field logic (priority: --name > config name > fallback)
         if args.name:
@@ -384,22 +375,13 @@ def handle_azure(args):
         assessment_type = require_env_int(
             "ESC_ASSESSMENT_TYPE", "assessment type (1 or 2)", {1, 2}
         )
-        try:
-            provider_details = _azure_provider_noninteractive(args)
-        except ConfigError:
-            sys.exit(codes.CONFIG)
+        provider_details = _azure_provider_noninteractive(args)
     elif args.cli:
-        try:
-            provider_details = _azure_provider_from_cli()
-        except ConfigError:
-            return
+        provider_details = _azure_provider_from_cli()
         exit_strategy, assessment_type = prompt_required_inputs()
     else:
         exit_strategy, assessment_type = prompt_required_inputs()
-        try:
-            provider_details = _azure_provider_from_prompt()
-        except ConfigError:
-            return
+        provider_details = _azure_provider_from_prompt()
 
     config = build_config(
         cloud_provider, exit_strategy, assessment_type, provider_details, args
@@ -864,17 +846,20 @@ def main():
         return
 
     # Dispatch based on provided arguments
-    if args.cloud_provider == "aws":
-        handle_aws(args)
-    elif args.cloud_provider == "azure":
-        handle_azure(args)
-    else:
-        console.print(
-            "[red]Invalid command. Use 'aws' or 'azure' as the first argument.[/red]"
-        )
-        console.print(
-            "[green]Run 'python3 main.py --help' for usage instructions.[/green]"
-        )
+    try:
+        if args.cloud_provider == "aws":
+            handle_aws(args)
+        elif args.cloud_provider == "azure":
+            handle_azure(args)
+        else:
+            console.print(
+                "[red]Invalid command. Use 'aws' or 'azure' as the first argument.[/red]"
+            )
+            console.print(
+                "[green]Run 'python3 main.py --help' for usage instructions.[/green]"
+            )
+    except ConfigError:
+        sys.exit(codes.CONFIG)
 
 
 if __name__ == "__main__":

tests/test_utils_and_main.py

@@ -58,7 +58,6 @@ def test_load_config_returns_none_for_invalid_json(self):
 
 
 def _base_patches():
-    """Patches that prevent real cloud/filesystem calls for all stage tests."""
     return [
         patch("main.console.print"),
         patch("main.print_step"),
@@ -93,7 +92,6 @@ def _base_patches():
 
 class RunAssessmentExitCodeTests(unittest.TestCase):
     def _run_with_patches(self, overrides: dict):
-        """Apply base patches, override specific ones, and run the assessment."""
         patches = _base_patches()
         for p in patches:
             p.start()
@@ -373,7 +371,6 @@ def test_dry_run_flag_passed_from_handle_aws(self):
 
 
 def _ni_aws_args(**kwargs):
-    """Build a Namespace that looks like 'aws --non-interactive' with optional overrides."""
     defaults = dict(
         config=None,
         profile=None,
@@ -387,7 +384,6 @@ def _ni_aws_args(**kwargs):
 
 
 def _ni_azure_args(**kwargs):
-    """Build a Namespace that looks like 'azure --non-interactive' with optional overrides."""
     defaults = dict(
         config=None,
         cli=False,
@@ -523,7 +519,7 @@ def test_uses_default_credential_when_client_secret_missing(self):
         self.assertEqual(config_arg["providerDetails"]["clientId"], "client-id-456")
         self.assertNotIn("clientSecret", config_arg["providerDetails"])
 
-    def test_missing_client_id_for_oidc_exits_config(self):
+    def test_missing_client_id_for_oidc_raises_config_error(self):
         env = {
             k: v
             for k, v in self._BASE_ENV.items()
@@ -533,9 +529,8 @@ def test_missing_client_id_for_oidc_exits_config(self):
             patch.dict(os.environ, env, clear=False),
             patch("main.console.print"),
         ):
-            with self.assertRaises(SystemExit) as ctx:
+            with self.assertRaises(main.ConfigError):
                 main.handle_azure(_ni_azure_args())
-        self.assertEqual(ctx.exception.code, codes.CONFIG)
 
     def test_missing_subscription_id_exits_config(self):
         env = {k: v for k, v in self._BASE_ENV.items() if k != "ESC_SUBSCRIPTION_ID"}
@@ -784,6 +779,44 @@ def test_handle_aws_passes_egress_flag(self):
         mock_run.assert_called_once_with(ANY, "aws", dry_run=False, egress=True)
 
 
+class MainExitCodeTests(unittest.TestCase):
+    def test_config_error_from_handler_exits_config(self):
+        with (
+            patch("main.initialize_dataset"),
+            patch("main.console.print"),
+            patch(
+                "main.parse_arguments",
+                return_value=Namespace(cloud_provider="aws"),
+            ),
+            patch("main.handle_aws", side_effect=main.ConfigError),
+        ):
+            with self.assertRaises(SystemExit) as ctx:
+                main.main()
+        self.assertEqual(ctx.exception.code, codes.CONFIG)
+
+    def test_bad_config_file_exits_config_end_to_end(self):
+        args = Namespace(
+            cloud_provider="aws",
+            config="/nonexistent/aws.json",
+            profile=None,
+            name=None,
+            non_interactive=False,
+            dry_run=False,
+            egress=False,
+        )
+        with (
+            patch("main.initialize_dataset"),
+            patch("main.console.print"),
+            patch("main.parse_arguments", return_value=args),
+            patch("main.load_config", return_value=None),
+            patch("main.run_assessment") as mock_run,
+        ):
+            with self.assertRaises(SystemExit) as ctx:
+                main.main()
+        self.assertEqual(ctx.exception.code, codes.CONFIG)
+        mock_run.assert_not_called()
+
+
 class ResolveModeEnvVarTests(unittest.TestCase):
     def test_host_and_key_env_override_config(self):
         fake_config = types.SimpleNamespace(HOST="config-host.io", KEY="config-key")