From e2e77029cefe01e529c9ed4d544ce48c842815a6 Mon Sep 17 00:00:00 2001
From: Kevin James <KevinJames@thekev.in>
Date: Thu, 29 Aug 2024 14:32:45 +0100
Subject: [PATCH] refactor: fix invalid escape sequences

Signed-off-by: Kevin James <KevinJames@thekev.in>
---
 .pre-commit-config.yaml                            |  6 ++++++
 backend/src/apiserver/visualization/types/tfdv.py  |  2 +-
 components/aws/athena/query/src/query.py           |  2 +-
 .../container/v1/aiplatform/remote_runner.py       |  2 +-
 .../v1/hyperparameter_tuning_job/component.py      |  2 +-
 proxy/get_proxy_url.py                             |  4 ++--
 sdk/python/kfp/compiler/compiler_test.py           | 12 ++++++------
 .../compiler/v2_compatible_compiler_test.py        |  6 +++---
 .../deprecated/components/type_annotation_utils.py |  2 +-
 .../deprecated/components_tests/test_python_op.py  |  2 +-
 sdk/python/kfp/deprecated/dsl/_ops_group.py        |  2 +-
 sdk/python/kfp/dsl/executor.py                     |  2 +-
 sdk/python/kfp/dsl/for_loop.py                     |  4 ++--
 sdk/python/kfp/dsl/placeholders_test.py            | 14 +++++++-------
 sdk/python/kfp/dsl/tasks_group_test.py             |  2 +-
 sdk/python/kfp/dsl/types/type_annotations.py       |  2 +-
 sdk/python/kfp/local/config_test.py                |  2 +-
 test/sample-test/sample_test_launcher.py           |  8 ++++----
 test/sample-test/unittests/utils_tests.py          | 10 +++++-----
 19 files changed, 46 insertions(+), 40 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 49e5c41120d..5e6a0f02599 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -15,6 +15,12 @@ repos:
       - id: double-quote-string-fixer
       - id: no-commit-to-branch
         args: [--branch, master]
+  - repo: https://github.com/PyCQA/flake8
+    rev: 7.1.1
+    hooks:
+    -   id: flake8
+        args:
+        - --select=W605
   # required formatting jobs (run these last)
 
   # add comment "noqa" to ignore an import that should not be removed
diff --git a/backend/src/apiserver/visualization/types/tfdv.py b/backend/src/apiserver/visualization/types/tfdv.py
index c8abfa849c3..262dc98a137 100644
--- a/backend/src/apiserver/visualization/types/tfdv.py
+++ b/backend/src/apiserver/visualization/types/tfdv.py
@@ -72,7 +72,7 @@ def get_statistics_html(
   # facets element and then remove it once we have appended the serialized proto
   # string to the element. We do this to avoid any collision of ids when
   # displaying multiple facets output in the notebook.
-  html_template = """<iframe id='facets-iframe' width="100%" height="500px"></iframe>
+  html_template = r"""<iframe id='facets-iframe' width="100%" height="500px"></iframe>
         <script>
         facets_iframe = document.getElementById('facets-iframe');
         facets_html = '<script src="https://cdnjs.cloudflare.com/ajax/libs/webcomponentsjs/1.3.3/webcomponents-lite.js"><\/script><link rel="import" href="https://raw.githubusercontent.com/PAIR-code/facets/master/facets-dist/facets-jupyter.html"><facets-overview proto-input="protostr"></facets-overview>';
diff --git a/components/aws/athena/query/src/query.py b/components/aws/athena/query/src/query.py
index 06a98dbc38c..db8d2963c34 100644
--- a/components/aws/athena/query/src/query.py
+++ b/components/aws/athena/query/src/query.py
@@ -63,7 +63,7 @@ def query(client, query, database, output, workgroup=None):
                     "OutputLocation"
                 ]
                 # could be multiple files?
-                filename = re.findall(".*\/(.*)", s3_path)[0]
+                filename = re.findall(r".*\/(.*)", s3_path)[0]
                 logging.info("S3 output file name %s", filename)
                 break
         time.sleep(5)
diff --git a/components/google-cloud/google_cloud_pipeline_components/container/v1/aiplatform/remote_runner.py b/components/google-cloud/google_cloud_pipeline_components/container/v1/aiplatform/remote_runner.py
index 8fdf6466c39..a00fb3ffa0f 100644
--- a/components/google-cloud/google_cloud_pipeline_components/container/v1/aiplatform/remote_runner.py
+++ b/components/google-cloud/google_cloud_pipeline_components/container/v1/aiplatform/remote_runner.py
@@ -88,7 +88,7 @@ def write_to_artifact(executor_input, text):
       # Add URI Prefix
       # "https://[location]-aiplatform.googleapis.com/API_VERSION/": For AI Platform resource names, current version is defined in AIPLATFORM_API_VERSION.
       if RESOURCE_NAME_PATTERN.match(text):
-        location = re.findall('locations/([\w\-]+)', text)[0]
+        location = re.findall(r'locations/([\w\-]+)', text)[0]
         uri_with_prefix = f'https://{location}-aiplatform.googleapis.com/{AIPLATFORM_API_VERSION}/{text}'
         metadata.update({'resourceName': text})
 
diff --git a/components/google-cloud/google_cloud_pipeline_components/v1/hyperparameter_tuning_job/component.py b/components/google-cloud/google_cloud_pipeline_components/v1/hyperparameter_tuning_job/component.py
index 34d04242390..24a5f2e0a64 100644
--- a/components/google-cloud/google_cloud_pipeline_components/v1/hyperparameter_tuning_job/component.py
+++ b/components/google-cloud/google_cloud_pipeline_components/v1/hyperparameter_tuning_job/component.py
@@ -42,7 +42,7 @@ def hyperparameter_tuning_job(
     project: str = _placeholders.PROJECT_ID_PLACEHOLDER,
 ):
   # fmt: off
-  """Creates a Vertex AI hyperparameter tuning job and waits for it to
+  r"""Creates a Vertex AI hyperparameter tuning job and waits for it to
   complete.
 
   See [more information](https://cloud.google.com/vertex-ai/docs/training/using-hyperparameter-tuning).
diff --git a/proxy/get_proxy_url.py b/proxy/get_proxy_url.py
index 619a2397336..11731718695 100644
--- a/proxy/get_proxy_url.py
+++ b/proxy/get_proxy_url.py
@@ -40,7 +40,7 @@ def urls_for_zone(zone, location_to_urls_map):
         ...
       }
   """
-    zone_match = re.match("((([a-z]+)-[a-z]+)\d+)-[a-z]", zone)
+    zone_match = re.match(r"((([a-z]+)-[a-z]+)\d+)-[a-z]", zone)
     if not zone_match:
         raise ValueError("Incorrect zone specified: {}".format(zone))
 
@@ -55,7 +55,7 @@ def urls_for_zone(zone, location_to_urls_map):
             url for url in location_to_urls_map[region] if url not in urls
         ])
 
-    region_regex = re.compile("([a-z]+-[a-z]+)\d+")
+    region_regex = re.compile(r"([a-z]+-[a-z]+)\d+")
     for location in location_to_urls_map:
         region_match = region_regex.match(location)
         if region_match and region_match.group(1) == approx_region:
diff --git a/sdk/python/kfp/compiler/compiler_test.py b/sdk/python/kfp/compiler/compiler_test.py
index d417d9eec19..b0741ab84da 100644
--- a/sdk/python/kfp/compiler/compiler_test.py
+++ b/sdk/python/kfp/compiler/compiler_test.py
@@ -4445,7 +4445,7 @@ def test_other_control_flow_instantiated_between_if_else_not_permitted(
             self):
         with self.assertRaisesRegex(
                 tasks_group.InvalidControlFlowException,
-                'dsl\.Else can only be used following an upstream dsl\.If or dsl\.Elif\.'
+                r'dsl\.Else can only be used following an upstream dsl\.If or dsl\.Elif\.'
         ):
 
             @dsl.pipeline
@@ -4908,7 +4908,7 @@ def flip_coin_pipeline(execute_pipeline: bool):
     def test_nested_under_condition_returned_raises(self):
         with self.assertRaisesRegex(
                 compiler_utils.InvalidTopologyException,
-                f'Pipeline outputs may only be returned from the top level of the pipeline function scope\. Got pipeline output dsl\.OneOf from within the control flow group dsl\.If\.'
+                r'Pipeline outputs may only be returned from the top level of the pipeline function scope\. Got pipeline output dsl\.OneOf from within the control flow group dsl\.If\.'
         ):
 
             @dsl.pipeline
@@ -4961,7 +4961,7 @@ def test_deeply_nested_returned_raises(self):
 
         with self.assertRaisesRegex(
                 compiler_utils.InvalidTopologyException,
-                f'Pipeline outputs may only be returned from the top level of the pipeline function scope\. Got pipeline output dsl\.OneOf from within the control flow group dsl\.ParallelFor\.'
+                r'Pipeline outputs may only be returned from the top level of the pipeline function scope\. Got pipeline output dsl\.OneOf from within the control flow group dsl\.ParallelFor\.'
         ):
 
             @dsl.pipeline
@@ -4983,7 +4983,7 @@ def test_consume_at_wrong_level(self):
 
         with self.assertRaisesRegex(
                 compiler_utils.InvalidTopologyException,
-                f'Illegal task dependency across DSL context managers\. A downstream task cannot depend on an upstream task within a dsl\.If context unless the downstream is within that context too\. Found task print-artifact which depends on upstream task condition-branches-5 within an uncommon dsl\.If context\.'
+                r'Illegal task dependency across DSL context managers\. A downstream task cannot depend on an upstream task within a dsl\.If context unless the downstream is within that context too\. Found task print-artifact which depends on upstream task condition-branches-5 within an uncommon dsl\.If context\.'
         ):
 
             @dsl.pipeline
@@ -5006,7 +5006,7 @@ def flip_coin_pipeline(execute_pipeline: bool):
     def test_return_at_wrong_level(self):
         with self.assertRaisesRegex(
                 compiler_utils.InvalidTopologyException,
-                f'Pipeline outputs may only be returned from the top level of the pipeline function scope\. Got pipeline output dsl\.OneOf from within the control flow group dsl\.If\.'
+                r'Pipeline outputs may only be returned from the top level of the pipeline function scope\. Got pipeline output dsl\.OneOf from within the control flow group dsl\.If\.'
         ):
 
             @dsl.pipeline
@@ -5143,7 +5143,7 @@ def roll_die_pipeline(
     def test_oneof_in_fstring(self):
         with self.assertRaisesRegex(
                 NotImplementedError,
-                f'dsl\.OneOf does not support string interpolation\.'):
+                r'dsl\.OneOf does not support string interpolation\.'):
 
             @dsl.pipeline
             def roll_die_pipeline():
diff --git a/sdk/python/kfp/deprecated/compiler/v2_compatible_compiler_test.py b/sdk/python/kfp/deprecated/compiler/v2_compatible_compiler_test.py
index 3d5177c72a9..482f7d7422c 100644
--- a/sdk/python/kfp/deprecated/compiler/v2_compatible_compiler_test.py
+++ b/sdk/python/kfp/deprecated/compiler/v2_compatible_compiler_test.py
@@ -88,7 +88,7 @@ def _load_compiled_template(filename: str) -> Dict:
 
                 if 'container' in template:
                     template['container'] = json.loads(
-                        re.sub("'kfp==(\d+).(\d+).(\d+)'", 'kfp',
+                        re.sub(r"'kfp==(\d+).(\d+).(\d+)'", 'kfp',
                                json.dumps(template['container'])))
 
             return workflow
@@ -154,8 +154,8 @@ def my_pipeline():
         with self.assertRaisesRegex(
                 RuntimeError,
                 'Constructing ContainerOp instances directly is deprecated and not '
-                'supported when compiling to v2 \(using v2 compiler or v1 compiler '
-                'with V2_COMPATIBLE or V2_ENGINE mode\).'):
+                r'supported when compiling to v2 \(using v2 compiler or v1 compiler '
+                r'with V2_COMPATIBLE or V2_ENGINE mode\)\.'):
             compiler.Compiler(
                 mode=v1dsl.PipelineExecutionMode.V2_COMPATIBLE).compile(
                     pipeline_func=my_pipeline, package_path='result.json')
diff --git a/sdk/python/kfp/deprecated/components/type_annotation_utils.py b/sdk/python/kfp/deprecated/components/type_annotation_utils.py
index b4b10daa6b4..0370937b0df 100644
--- a/sdk/python/kfp/deprecated/components/type_annotation_utils.py
+++ b/sdk/python/kfp/deprecated/components/type_annotation_utils.py
@@ -58,7 +58,7 @@ def get_short_type_name(type_name: str) -> str:
     Returns:
       The short form type name or the original name if pattern doesn't match.
     """
-    match = re.match('(typing\.)?(?P<type>\w+)(?:\[.+\])?', type_name)
+    match = re.match(r'(typing\.)?(?P<type>\w+)(?:\[.+\])?', type_name)
     if match:
         return match.group('type')
     else:
diff --git a/sdk/python/kfp/deprecated/components_tests/test_python_op.py b/sdk/python/kfp/deprecated/components_tests/test_python_op.py
index 4463485b396..a116ad645ac 100644
--- a/sdk/python/kfp/deprecated/components_tests/test_python_op.py
+++ b/sdk/python/kfp/deprecated/components_tests/test_python_op.py
@@ -1224,7 +1224,7 @@ def my_func(args: Sequence[int]):
         task_factory = comp.create_component_from_func(my_func)
         with self.assertRaisesRegex(
                 TypeError,
-                'There are no registered serializers for type "(typing.)?Sequence(\[int\])?"'
+                r'There are no registered serializers for type "(typing.)?Sequence(\[int\])?"'
         ):
             self.helper_test_component_using_local_call(
                 task_factory, arguments={'args': [1, 2, 3]})
diff --git a/sdk/python/kfp/deprecated/dsl/_ops_group.py b/sdk/python/kfp/deprecated/dsl/_ops_group.py
index 709d02e90ad..65fa90c4338 100644
--- a/sdk/python/kfp/deprecated/dsl/_ops_group.py
+++ b/sdk/python/kfp/deprecated/dsl/_ops_group.py
@@ -68,7 +68,7 @@ def _get_matching_opsgroup_already_in_pipeline(group_type, name):
         if name is None:
             return None
         name_pattern = '^' + (group_type + '-' + name + '-').replace(
-            '_', '-') + '[\d]+$'
+            '_', '-') + r'[\d]+$'
         for ops_group_already_in_pipeline in _pipeline.Pipeline.get_default_pipeline(
         ).groups:
             import re
diff --git a/sdk/python/kfp/dsl/executor.py b/sdk/python/kfp/dsl/executor.py
index 272c3c1f611..0115e2c376c 100644
--- a/sdk/python/kfp/dsl/executor.py
+++ b/sdk/python/kfp/dsl/executor.py
@@ -400,7 +400,7 @@ def get_short_type_name(type_name: str) -> str:
     Returns:
       The short form type name or the original name if pattern doesn't match.
     """
-    match = re.match('(typing\.)?(?P<type>\w+)(?:\[.+\])?', type_name)
+    match = re.match(r'(typing\.)?(?P<type>\w+)(?:\[.+\])?', type_name)
     return match['type'] if match else type_name
 
 
diff --git a/sdk/python/kfp/dsl/for_loop.py b/sdk/python/kfp/dsl/for_loop.py
index 9c4b8f69585..e7e8e552310 100644
--- a/sdk/python/kfp/dsl/for_loop.py
+++ b/sdk/python/kfp/dsl/for_loop.py
@@ -44,7 +44,7 @@ def _get_loop_item_type(type_name: str) -> Optional[str]:
     Returns:
         The collection item type or None if no match found.
     """
-    match = re.match('(typing\.)?(?:\w+)(?:\[(?P<item_type>.+)\])', type_name)
+    match = re.match(r'(typing\.)?(?:\w+)(?:\[(?P<item_type>.+)\])', type_name)
     return match['item_type'].lstrip().rstrip() if match else None
 
 
@@ -64,7 +64,7 @@ def _get_subvar_type(type_name: str) -> Optional[str]:
         The dictionary value type or None if no match found.
     """
     match = re.match(
-        '(typing\.)?(?:\w+)(?:\[\s*(?:\w+)\s*,\s*(?P<value_type>.+)\])',
+        r'(typing\.)?(?:\w+)(?:\[\s*(?:\w+)\s*,\s*(?P<value_type>.+)\])',
         type_name)
     return match['value_type'].lstrip().rstrip() if match else None
 
diff --git a/sdk/python/kfp/dsl/placeholders_test.py b/sdk/python/kfp/dsl/placeholders_test.py
index a6310ed5b5b..902a2af41bf 100644
--- a/sdk/python/kfp/dsl/placeholders_test.py
+++ b/sdk/python/kfp/dsl/placeholders_test.py
@@ -299,7 +299,7 @@ def test_valid(self, placeholder_obj: placeholders.IfPresentPlaceholder,
     def test_only_single_element_ifpresent_inside_concat_outer(self):
         with self.assertRaisesRegex(
                 ValueError,
-                f'Please use a single element for `then` and `else_` only\.'):
+                r'Please use a single element for `then` and `else_` only\.'):
             placeholders.ConcatPlaceholder([
                 'b',
                 placeholders.IfPresentPlaceholder(
@@ -309,7 +309,7 @@ def test_only_single_element_ifpresent_inside_concat_outer(self):
     def test_only_single_element_ifpresent_inside_concat_recursive(self):
         with self.assertRaisesRegex(
                 ValueError,
-                f'Please use a single element for `then` and `else_` only\.'):
+                r'Please use a single element for `then` and `else_` only\.'):
             placeholders.ConcatPlaceholder([
                 'a',
                 placeholders.ConcatPlaceholder([
@@ -323,7 +323,7 @@ def test_only_single_element_ifpresent_inside_concat_recursive(self):
 
         with self.assertRaisesRegex(
                 ValueError,
-                f'Please use a single element for `then` and `else_` only\.'):
+                r'Please use a single element for `then` and `else_` only\.'):
             placeholders.ConcatPlaceholder([
                 'a',
                 placeholders.ConcatPlaceholder([
@@ -341,7 +341,7 @@ def test_only_single_element_ifpresent_inside_concat_recursive(self):
     def test_only_single_element_in_nested_ifpresent_inside_concat(self):
         with self.assertRaisesRegex(
                 ValueError,
-                f'Please use a single element for `then` and `else_` only\.'):
+                r'Please use a single element for `then` and `else_` only\.'):
             dsl.ConcatPlaceholder([
                 'my-prefix-',
                 dsl.IfPresentPlaceholder(
@@ -357,7 +357,7 @@ def test_recursive_nested_placeholder_validation_does_not_exit_when_first_valid_
             self):
         with self.assertRaisesRegex(
                 ValueError,
-                f'Please use a single element for `then` and `else_` only\.'):
+                r'Please use a single element for `then` and `else_` only\.'):
             dsl.ConcatPlaceholder([
                 'my-prefix-',
                 dsl.IfPresentPlaceholder(
@@ -371,7 +371,7 @@ def test_only_single_element_in_nested_ifpresent_inside_concat_with_outer_ifpres
             self):
         with self.assertRaisesRegex(
                 ValueError,
-                f'Please use a single element for `then` and `else_` only\.'):
+                r'Please use a single element for `then` and `else_` only\.'):
             dsl.IfPresentPlaceholder(
                 input_name='input_1',
                 then=dsl.ConcatPlaceholder([
@@ -386,7 +386,7 @@ def test_only_single_element_in_nested_ifpresent_inside_concat_with_outer_ifpres
     def test_valid_then_but_invalid_else(self):
         with self.assertRaisesRegex(
                 ValueError,
-                f'Please use a single element for `then` and `else_` only\.'):
+                r'Please use a single element for `then` and `else_` only\.'):
             dsl.ConcatPlaceholder([
                 'my-prefix-',
                 dsl.IfPresentPlaceholder(
diff --git a/sdk/python/kfp/dsl/tasks_group_test.py b/sdk/python/kfp/dsl/tasks_group_test.py
index 92bcf505887..0d298e09a84 100644
--- a/sdk/python/kfp/dsl/tasks_group_test.py
+++ b/sdk/python/kfp/dsl/tasks_group_test.py
@@ -74,7 +74,7 @@ def foo() -> str:
         def my_pipeline(string: str):
             with self.assertWarnsRegex(
                     DeprecationWarning,
-                    'dsl\.Condition is deprecated\. Please use dsl\.If instead\.'
+                    r'dsl\.Condition is deprecated\. Please use dsl\.If instead\.'
             ):
                 with dsl.Condition(string == 'text'):
                     foo()
diff --git a/sdk/python/kfp/dsl/types/type_annotations.py b/sdk/python/kfp/dsl/types/type_annotations.py
index aa5776b9988..4f1c99ba92c 100644
--- a/sdk/python/kfp/dsl/types/type_annotations.py
+++ b/sdk/python/kfp/dsl/types/type_annotations.py
@@ -227,7 +227,7 @@ def get_short_type_name(type_name: str) -> str:
     Returns:
       The short form type name or the original name if pattern doesn't match.
     """
-    match = re.match('(typing\.)?(?P<type>\w+)(?:\[.+\])?', type_name)
+    match = re.match(r'(typing\.)?(?P<type>\w+)(?:\[.+\])?', type_name)
     return match['type'] if match else type_name
 
 
diff --git a/sdk/python/kfp/local/config_test.py b/sdk/python/kfp/local/config_test.py
index 3e1ebb26c59..ad71cf5a6ac 100755
--- a/sdk/python/kfp/local/config_test.py
+++ b/sdk/python/kfp/local/config_test.py
@@ -70,7 +70,7 @@ def test_validate_success(self):
     def test_validate_fail(self):
         with self.assertRaisesRegex(
                 RuntimeError,
-                f"Local environment not initialized. Please run 'kfp\.local\.init\(\)' before executing tasks locally\."
+                r"Local environment not initialized. Please run 'kfp\.local\.init\(\)' before executing tasks locally\."
         ):
             config.LocalExecutionConfig.validate()
 
diff --git a/test/sample-test/sample_test_launcher.py b/test/sample-test/sample_test_launcher.py
index 3de1c06d7be..dabbc790921 100644
--- a/test/sample-test/sample_test_launcher.py
+++ b/test/sample-test/sample_test_launcher.py
@@ -92,7 +92,7 @@ def _compile(self):
         for file in list_of_files:
             # matching by .py or .ipynb, there will be yaml ( compiled ) files in the folder.
             # if you rerun the test suite twice, the test suite will fail
-            m = re.match(self._test_name + '\.(py|ipynb)$', file)
+            m = re.match(self._test_name + r'\.(py|ipynb)$', file)
             if m:
                 file_name, ext_name = os.path.splitext(file)
                 if self._is_notebook is not None:
@@ -242,14 +242,14 @@ def __init__(self,
     def _injection(self):
         """Sample-specific image injection into yaml file."""
         subs = {  # Tag can look like 1.0.0-rc.3, so we need both "-" and "." in the regex.
-            'gcr\.io/ml-pipeline/ml-pipeline/ml-pipeline-local-confusion-matrix:(\w+|[.-])+':
+            r'gcr\.io/ml-pipeline/ml-pipeline/ml-pipeline-local-confusion-matrix:(\w+|[.-])+':
                 self._local_confusionmatrix_image,
-            'gcr\.io/ml-pipeline/ml-pipeline/ml-pipeline-local-roc:(\w+|[.-])+':
+            r'gcr\.io/ml-pipeline/ml-pipeline/ml-pipeline-local-roc:(\w+|[.-])+':
                 self._local_roc_image
         }
         if self._test_name == 'xgboost_training_cm':
             subs.update({
-                'gcr\.io/ml-pipeline/ml-pipeline-gcp:(\w|[.-])+':
+                r'gcr\.io/ml-pipeline/ml-pipeline-gcp:(\w|[.-])+':
                     self._dataproc_gcp_image
             })
 
diff --git a/test/sample-test/unittests/utils_tests.py b/test/sample-test/unittests/utils_tests.py
index ab5fb6b89e6..0a026ee566f 100644
--- a/test/sample-test/unittests/utils_tests.py
+++ b/test/sample-test/unittests/utils_tests.py
@@ -48,10 +48,10 @@ def tearDown(self) -> None:
   def test_file_injection(self):
     """Test file_injection function."""
     subs = {
-        'gcr\.io/ml-pipeline/ml-pipeline-local-confusion-matrix:\w+':'gcr.io/ml-pipeline/LOCAL_CONFUSION_MATRIX_IMAGE',
-        'gcr\.io/ml-pipeline/ml-pipeline-dataproc-analyze:\w+':'gcr.io/ml-pipeline/DATAPROC_ANALYZE_IMAGE',
-        'gcr\.io/ml-pipeline/ml-pipeline-dataproc-create-cluster:\w+':'gcr.io/ml-pipeline/DATAPROC_CREATE_IMAGE',
-        'gcr\.io/ml-pipeline/ml-pipeline-dataproc-delete-cluster:\w+':'gcr.io/ml-pipeline/DATAPROC_DELETE_IMAGE',
+        r'gcr\.io/ml-pipeline/ml-pipeline-local-confusion-matrix:\w+':'gcr.io/ml-pipeline/LOCAL_CONFUSION_MATRIX_IMAGE',
+        r'gcr\.io/ml-pipeline/ml-pipeline-dataproc-analyze:\w+':'gcr.io/ml-pipeline/DATAPROC_ANALYZE_IMAGE',
+        r'gcr\.io/ml-pipeline/ml-pipeline-dataproc-create-cluster:\w+':'gcr.io/ml-pipeline/DATAPROC_CREATE_IMAGE',
+        r'gcr\.io/ml-pipeline/ml-pipeline-dataproc-delete-cluster:\w+':'gcr.io/ml-pipeline/DATAPROC_DELETE_IMAGE',
     }
     utils.file_injection(
         os.path.join(_WORK_DIR, 'test_file_injection.yaml'),
@@ -64,4 +64,4 @@ def test_file_injection(self):
     with open(os.path.join(_WORK_DIR,
                            'test_file_injection.yaml'), 'r') as f:
       injected = yaml.safe_load(f)
-    self.assertEqual(golden, injected)
\ No newline at end of file
+    self.assertEqual(golden, injected)