offloading error checking and updating tests

src/confcom/azext_confcom/_params.py

@@ -5,6 +5,15 @@
 # pylint: disable=line-too-long
 
 from knack.arguments import CLIArgumentType
+from azext_confcom._validators import (
+    validate_params_file,
+    validate_diff,
+    validate_aci_source,
+    validate_print_format,
+    validate_save_to_file,
+    validate_faster_hashing,
+    validate_katapolicygen_input,
+)
 
 
 def load_arguments(self, _):
@@ -27,30 +36,35 @@ def load_arguments(self, _):
             options_list=("--input", "-i"),
             required=False,
             help="Input JSON config file",
+            validator=validate_aci_source
         )
         c.argument(
             "arm_template",
             options_list=("--template-file", "-a"),
             required=False,
             help="ARM template file",
+            validator=validate_aci_source
         )
         c.argument(
             "arm_template_parameters",
             options_list=("--parameters", "-p"),
             required=False,
             help="ARM template parameters",
+            validator=validate_params_file
         )
         c.argument(
             "virtual_node_yaml_path",
             options_list=("--virtual-node-yaml"),
             required=False,
             help="Virtual node YAML file",
+            validator=validate_aci_source
         )
         c.argument(
             "image_name",
             options_list=("--image",),
             required=False,
             help="Image Name",
+            validator=validate_aci_source
         )
         c.argument(
             "tar_mapping_location",
@@ -87,6 +101,7 @@ def load_arguments(self, _):
             options_list=("--diff", "-d"),
             required=False,
             help="Compare the CCE Policy field in the ARM Template to the containers in the ARM Template and make sure they are compatible",
+            validator=validate_diff
         )
         c.argument(
             "validate_sidecar",
@@ -107,31 +122,36 @@ def load_arguments(self, _):
             required=False,
             action="store_true",
             help="Output policy in clear text compact JSON instead of default base64 format",
+            validator=validate_print_format,
         )
         c.argument(
             "outraw_pretty_print",
             options_list=("--outraw-pretty-print"),
             required=False,
             action="store_true",
             help="Output policy in clear text and pretty print format",
+            validator=validate_print_format,
         )
         c.argument(
             "save_to_file",
             options_list=("--save-to-file", "-s"),
             required=False,
             help="Save output policy to given file path",
+            validator=validate_save_to_file,
         )
         c.argument(
             "print_policy_to_terminal",
             options_list=("--print-policy"),
             required=False,
             help="Print the generated policy in the terminal",
+            validator=validate_print_format,
         )
         c.argument(
             "faster_hashing",
             options_list=("--faster-hashing"),
             required=False,
             help="Use buffered image reader for dmverity hashing. This will speed up the hashing process but use much more memory.",
+            validator=validate_faster_hashing,
         )
 
     with self.argument_context("confcom katapolicygen") as c:
@@ -140,58 +160,68 @@ def load_arguments(self, _):
             options_list=("--yaml", "-y"),
             required=False,
             help="Input YAML config file",
+            validator=validate_katapolicygen_input,
         )
         c.argument(
             "outraw",
             options_list=("--outraw"),
             required=False,
             help="Print the generated policy in the terminal in Rego format",
+            validator=validate_katapolicygen_input,
         )
         c.argument(
             "print_policy",
             options_list=("--print-policy"),
             required=False,
             help="Print the generated policy in the terminal in base64",
+            validator=validate_katapolicygen_input,
         )
         c.argument(
             "config_map_file",
             options_list=("--config-map-file", "-c"),
             required=False,
             help="Config map file",
+            validator=validate_katapolicygen_input,
         )
         c.argument(
             "use_cached_files",
             options_list=("--use-cached-files", "-u"),
             required=False,
             help="Use cached files",
+            validator=validate_katapolicygen_input,
         )
         c.argument(
             "settings_file_name",
             options_list=("--settings-file-name", "-j"),
             required=False,
             help="Path for custom settings file",
+            validator=validate_katapolicygen_input,
         )
         c.argument(
             "rules_file_name",
             options_list=("--rules-file-name", "-p"),
             required=False,
             help="Path for custom rules file",
+            validator=validate_katapolicygen_input,
         )
         c.argument(
             "print_version",
             options_list=("--print-version", "-v"),
             required=False,
             help="Print the version of the genpolicy tool",
+            validator=validate_katapolicygen_input,
         )
         c.argument(
             "containerd_pull",
             options_list=("--containerd-pull", "-d"),
             required=False,
             help="Use containerd to pull the image",
+            validator=validate_katapolicygen_input,
         )
         c.argument(
             "containerd_socket_path",
             options_list=("--containerd-socket-path"),
             required=False,
             help="Path to containerd socket if not using the default",
+            validator=validate_katapolicygen_input,
         )

src/confcom/azext_confcom/_validators.py

@@ -3,20 +3,48 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
+from knack.util import CLIError
 
-def example_name_or_id_validator(cmd, namespace):
-    # Example of a storage account name or ID validator.
-    # pylint: disable=line-too-long
-    # See: https://github.com/Azure/azure-cli/blob/dev/doc/authoring_command_modules/authoring_commands.md#supporting-name-or-id-parameters
-    from azure.cli.core.commands.client_factory import get_subscription_id
-    from msrestazure.tools import is_valid_resource_id, resource_id
-
-    if namespace.storage_account:
-        if not is_valid_resource_id(namespace.RESOURCE):
-            namespace.storage_account = resource_id(
-                subscription=get_subscription_id(cmd.cli_ctx),
-                resource_group=namespace.resource_group_name,
-                namespace="Microsoft.Storage",
-                type="storageAccounts",
-                name=namespace.storage_account,
-            )
+
+def validate_params_file(namespace):
+    if namespace.arm_template_parameters and not namespace.arm_template:
+        raise CLIError(
+            "Can only use ARM Template Parameters if ARM Template is also present"
+        )
+
+
+def validate_diff(namespace):
+    if (namespace.diff and namespace.input_path) or (namespace.diff and namespace.image_name):
+        raise CLIError("Can only diff CCE policy from ARM Template or YAML File")
+
+
+def validate_print_format(namespace):
+    if sum(map(bool, [namespace.print_policy_to_terminal, namespace.outraw, namespace.outraw_pretty_print])) > 1:
+        raise CLIError("Can only print in one format at a time")
+
+
+def validate_aci_source(namespace):
+    if sum(map(bool, [
+        namespace.input_path,
+        namespace.arm_template,
+        namespace.image_name,
+        namespace.virtual_node_yaml_path
+    ])) != 1:
+        raise CLIError("Can only generate CCE policy from one source at a time")
+
+
+def validate_faster_hashing(namespace):
+    if namespace.faster_hashing and namespace.tar_mapping_location:
+        raise CLIError("Cannot use --faster-hashing with --tar")
+
+
+def validate_save_to_file(namespace):
+    if namespace.save_to_file and namespace.arm_template and not (
+        namespace.print_policy_to_terminal or namespace.outraw or namespace.outraw_pretty_print
+    ):
+        raise CLIError("Must print policy to terminal when saving to file")
+
+
+def validate_katapolicygen_input(namespace):
+    if namespace.yaml_path and not namespace.print_version:
+        raise CLIError("Either --yaml-path or --print-version is required")

src/confcom/azext_confcom/custom.py

@@ -57,21 +57,6 @@ def acipolicygen_confcom(
     faster_hashing: bool = False,
 ):
 
-    if sum(map(bool, [input_path, arm_template, image_name, virtual_node_yaml_path])) != 1:
-        error_out("Can only generate CCE policy from one source at a time")
-    if sum(map(bool, [print_policy_to_terminal, outraw, outraw_pretty_print])) > 1:
-        error_out("Can only print in one format at a time")
-    elif (diff and input_path) or (diff and image_name):
-        error_out("Can only diff CCE policy from ARM Template or YAML File")
-    elif arm_template_parameters and not arm_template:
-        error_out(
-            "Can only use ARM Template Parameters if ARM Template is also present"
-        )
-    elif save_to_file and arm_template and not (print_policy_to_terminal or outraw or outraw_pretty_print):
-        error_out("Must print policy to terminal when saving to file")
-    elif faster_hashing and tar_mapping_location:
-        error_out("Cannot use --faster-hashing with --tar")
-
     if print_existing_policy or outraw or outraw_pretty_print:
         logger.warning(
             "%s %s %s %s %s",
@@ -215,9 +200,6 @@ def katapolicygen_confcom(
 ):
     kata_proxy = KataPolicyGenProxy()
 
-    if not (yaml_path or print_version):
-        error_out("Either --yaml-path or --print-version is required")
-
     output = kata_proxy.kata_genpolicy(
         yaml_path,
         config_map_file=config_map_file,
@@ -327,8 +309,3 @@ def get_output_type(outraw, outraw_pretty_print):
     elif outraw_pretty_print:
         output_type = security_policy.OutputType.PRETTY_PRINT
     return output_type
-
-
-def error_out(error_string):
-    logger.error(error_string)
-    sys.exit(1)

src/confcom/azext_confcom/tests/latest/test_confcom_arm.py

@@ -1362,15 +1362,15 @@ def test_arm_template_with_parameter_file_clean_room(self):
             except:
                 # do nothing
                 pass
-            regular_image = load_policy_from_arm_template_str(
-                custom_arm_json_default_value, ""
-            )
-            regular_image[0].populate_policy_content_for_all_images()
+        regular_image = load_policy_from_arm_template_str(
+            custom_arm_json_default_value, ""
+        )
+        regular_image[0].populate_policy_content_for_all_images()
             # create and tag same image to the new name to see if docker will error out that the image is not in a remote repo
-            new_repo = "fakerepo.microsoft.com"
-            new_image_name = "azure-functions"
-            new_tag = "fake-tag"
-
+        new_repo = "fakerepo.microsoft.com"
+        new_image_name = "azure-functions"
+        new_tag = "fake-tag"
+        with DockerClient() as client:
             image = client.images.get(original_image)
             try:
                 client.images.remove(new_repo + "/" + new_image_name + ":" + new_tag)

src/confcom/azext_confcom/tests/latest/test_confcom_image.py

@@ -14,6 +14,7 @@
     load_policy_from_image_name,
     load_policy_from_str,
 )
+from  azext_confcom.template_util import DockerClient
 import azext_confcom.config as config
 
 TEST_DIR = os.path.abspath(os.path.join(os.path.abspath(__file__), ".."))
@@ -99,35 +100,34 @@ def test_invalid_image_policy(self):
 
 class PolicyGeneratingImageCleanRoom(unittest.TestCase):
     def test_clean_room_policy(self):
-        client = docker.from_env()
-        original_image = (
-            "mcr.microsoft.com/aci/atlas-mount-azure-file-volume:master_20201210.2"
-        )
-        try:
-            client.images.remove(original_image)
-        except:
-            # do nothing
-            pass
+        with DockerClient() as client:
+            original_image = (
+                "mcr.microsoft.com/aci/atlas-mount-azure-file-volume:master_20201210.2"
+            )
+            try:
+                client.images.remove(original_image)
+            except:
+                # do nothing
+                pass
         regular_image = load_policy_from_image_name(original_image)
         regular_image.populate_policy_content_for_all_images(individual_image=True)
         # create and tag same image to the new name to see if docker will error out that the image is not in a remote repo
         new_repo = "mcr.microsoft.com"
         new_image_name = "aci/atlas-mount-azure-file-volume"
         new_tag = "fake-tag"
-
-        image = client.images.get(original_image)
-        try:
-            client.images.remove(new_repo + "/" + new_image_name + ":" + new_tag)
-        except:
-            # do nothing
-            pass
-        image.tag(new_repo + "/" + new_image_name, tag=new_tag)
-        try:
-            client.images.remove(original_image)
-        except:
-            # do nothing
-            pass
-        client.close()
+        with DockerClient() as client:
+            image = client.images.get(original_image)
+            try:
+                client.images.remove(new_repo + "/" + new_image_name + ":" + new_tag)
+            except:
+                # do nothing
+                pass
+            image.tag(new_repo + "/" + new_image_name, tag=new_tag)
+            try:
+                client.images.remove(original_image)
+            except:
+                # do nothing
+                pass
 
         policy = load_policy_from_image_name(
             new_repo + "/" + new_image_name + ":" + new_tag