diff --git a/sdk/python/tests/compiler/compiler_tests.py b/sdk/python/tests/compiler/compiler_tests.py
index 181e5b8d950..b160959c190 100644
--- a/sdk/python/tests/compiler/compiler_tests.py
+++ b/sdk/python/tests/compiler/compiler_tests.py
@@ -31,6 +31,14 @@
 from kubernetes.client import V1Toleration
 
 
+def some_op():
+  return dsl.ContainerOp(
+      name='sleep',
+      image='busybox',
+      command=['sleep 1'],
+  )
+
+
 class TestCompiler(unittest.TestCase):
   # Define the places of samples covered by unit tests.
   core_sample_path = os.path.join(os.path.dirname(__file__), '..', '..', '..',
@@ -337,7 +345,16 @@ def test_py_volume(self):
 
   def test_py_retry(self):
     """Test retry functionality."""
-    self._test_py_compile_yaml('retry')
+    number_of_retries = 137
+    def my_pipeline():
+      some_op().set_retry(number_of_retries)
+
+    workflow = kfp.compiler.Compiler()._compile(my_pipeline)
+    name_to_template = {template['name']: template for template in workflow['spec']['templates']}
+    main_dag_tasks = name_to_template[workflow['spec']['entrypoint']]['dag']['tasks']
+    template = name_to_template[main_dag_tasks[0]['template']]
+
+    self.assertEqual(template['retryStrategy']['limit'], number_of_retries)
 
   def test_py_image_pull_secrets(self):
     """Test pipeline imagepullsecret."""
diff --git a/sdk/python/tests/compiler/testdata/retry.py b/sdk/python/tests/compiler/testdata/retry.py
deleted file mode 100755
index 030fc1ddc49..00000000000
--- a/sdk/python/tests/compiler/testdata/retry.py
+++ /dev/null
@@ -1,42 +0,0 @@
-#!/usr/bin/env python3
-# Copyright 2018 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
-import kfp.dsl as dsl
-
-
-class RandomFailure1Op(dsl.ContainerOp):
-  """A component that fails randomly."""
-
-  def __init__(self, exit_codes):
-    super(RandomFailure1Op, self).__init__(
-      name='random_failure',
-      image='python:alpine3.6',
-      command=['python', '-c'],
-      arguments=["import random; import sys; exit_code = random.choice([%s]); print(exit_code); sys.exit(exit_code)" % exit_codes])
-
-
-@dsl.pipeline(
-  name='pipeline includes two steps which fail randomly.',
-  description='shows how to use ContainerOp set_retry().'
-)
-def retry_sample_pipeline():
-  op1 = RandomFailure1Op('0,1,2,3').set_retry(100)
-  op2 = RandomFailure1Op('0,1').set_retry(50)
-
-
-if __name__ == '__main__':
-  import kfp.compiler as compiler
-  compiler.Compiler().compile(retry_sample_pipeline, __file__ + '.tar.gz')
diff --git a/sdk/python/tests/compiler/testdata/retry.yaml b/sdk/python/tests/compiler/testdata/retry.yaml
deleted file mode 100644
index 5c40177f48d..00000000000
--- a/sdk/python/tests/compiler/testdata/retry.yaml
+++ /dev/null
@@ -1,59 +0,0 @@
-apiVersion: argoproj.io/v1alpha1
-kind: Workflow
-metadata:
-  annotations:
-    pipelines.kubeflow.org/pipeline_spec: '{"description": "shows how to use ContainerOp
-      set_retry().", "inputs": [], "name": "pipeline includes two steps which fail
-      randomly."}'
-  generateName: pipeline-includes-two-steps-which-fail-randomly-
-spec:
-  arguments:
-    parameters: []
-  entrypoint: pipeline-includes-two-steps-which-fail-randomly
-  serviceAccountName: pipeline-runner
-  templates:
-  - dag:
-      tasks:
-      - name: random-failure
-        template: random-failure
-      - name: random-failure-2
-        template: random-failure-2
-    name: pipeline-includes-two-steps-which-fail-randomly
-  - container:
-      args:
-      - import random; import sys; exit_code = random.choice([0,1,2,3]); print(exit_code);
-        sys.exit(exit_code)
-      command:
-      - python
-      - -c
-      image: python:alpine3.6
-    name: random-failure
-    outputs:
-      artifacts:
-      - name: mlpipeline-ui-metadata
-        path: /mlpipeline-ui-metadata.json
-        optional: true
-      - name: mlpipeline-metrics
-        path: /mlpipeline-metrics.json
-        optional: true
-    retryStrategy:
-      limit: 100
-  - container:
-      args:
-      - import random; import sys; exit_code = random.choice([0,1]); print(exit_code);
-        sys.exit(exit_code)
-      command:
-      - python
-      - -c
-      image: python:alpine3.6
-    name: random-failure-2
-    outputs:
-      artifacts:
-      - name: mlpipeline-ui-metadata
-        path: /mlpipeline-ui-metadata.json
-        optional: true
-      - name: mlpipeline-metrics
-        path: /mlpipeline-metrics.json
-        optional: true
-    retryStrategy:
-      limit: 50