Add dataproc network tag for smoke test

custom_image_utils/args_parser.py

@@ -47,7 +47,7 @@ def _full_image_family_uri_regex_type(s):
   """Check if the partial image family uri string matches regex."""
   if not _FULL_IMAGE_FAMILY_URI.match(s):
     raise argparse.ArgumentTypeError("Invalid image family URI: {}.".format(s))
-  return s  
+  return s
 
 def parse_args(args):
   """Parses command-line arguments."""
@@ -75,7 +75,7 @@ def parse_args(args):
       "--base-image-family",
       type=_full_image_family_uri_regex_type,
       help="""The source image family URI. The latest non-depracated image associated with the family will be used.
-      """)      
+      """)
   required_args.add_argument(
       "--customization-script",
       type=str,
@@ -154,6 +154,13 @@ def parse_args(args):
       that builds the custom image. A full subnetwork URL is required.
       Default subnetwork is None. For shared VPC only provide this parameter and
       do not use the --network argument.""")
+  parser.add_argument(
+      "--smoke-test-tags",
+      type=str,
+      required=False,
+      default="",
+      help="""(Optional) Network tags used to launch the Dataproc cluster used
+      for the smoke test.""")
   parser.add_argument(
       "--no-external-ip",
       action="store_true",

custom_image_utils/smoke_test_runner.py

@@ -24,7 +24,7 @@
 _LOG.setLevel(logging.INFO)
 
 def _create_workflow_template(workflow_name, image_name, project_id, zone, region,
-                              network, subnet, no_external_ip):
+                              network, subnet, no_external_ip, tags):
   """Create a Dataproc workflow template for testing."""
   create_command = [
       "gcloud", "dataproc", "workflow-templates", "create",
@@ -41,6 +41,8 @@ def _create_workflow_template(workflow_name, image_name, project_id, zone, regio
     set_cluster_command.extend(["--subnet", subnet])
   if no_external_ip:
     set_cluster_command.extend(["--no-address"])
+  if tags:
+    set_cluster_command.extend(["--tags", tags])
   add_job_command = [
       "gcloud", "dataproc", "workflow-templates", "add-job", "spark",
       "--workflow-template", workflow_name, "--project", project_id, "--region", region,
@@ -92,7 +94,7 @@ def _delete_workflow_template(workflow_name, project_id, region):
     raise RuntimeError("Error deleting workfloe template %s.", workflow_name)
 
 
-def _verify_custom_image(image_name, project_id, zone, network, subnetwork, no_external_ip):
+def _verify_custom_image(image_name, project_id, zone, network, subnetwork, no_external_ip, tags):
   """Verifies if custom image works with Dataproc."""
   region = zone[:-2]
   date = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
@@ -103,7 +105,7 @@ def _verify_custom_image(image_name, project_id, zone, network, subnetwork, no_e
     _LOG.info("Creating Dataproc workflow-template %s with image %s...",
               workflow_name, image_name)
     _create_workflow_template(workflow_name, image_name, project_id, zone, region,
-                              network, subnetwork, no_external_ip)
+                              network, subnetwork, no_external_ip, tags)
     _LOG.info(
         "Successfully created Dataproc workflow-template %s with image %s...",
         workflow_name, image_name)
@@ -133,7 +135,8 @@ def run(args):
     if not args.no_smoke_test:
       _LOG.info("Verifying the custom image...")
       _verify_custom_image(args.image_name, args.project_id, args.zone,
-                           args.network, args.subnetwork, args.no_external_ip)
+                           args.network, args.subnetwork, args.no_external_ip,
+                           args.smoke_test_tags)
       _LOG.info("Successfully verified the custom image...")
   else:
     _LOG.info("Skip running smoke test (dry run).")