diff --git a/main.py b/main.py index f01b264..943c0d5 100644 --- a/main.py +++ b/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__": diff --git a/tests/test_utils_and_main.py b/tests/test_utils_and_main.py index dacc0c7..eb16c99 100644 --- a/tests/test_utils_and_main.py +++ b/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")