Added test for op.set_display_name

sdk/python/tests/compiler/compiler_tests.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import kfp
 import kfp.compiler as compiler
 import kfp.dsl as dsl
 import os
@@ -497,6 +498,27 @@ def test_tolerations(self):
 
     self._test_op_to_template_yaml(op1, file_base_name='tolerations')
 
+  def test_set_display_name(self):
+    """Test a pipeline with a customized task names."""
+
+    import kfp
+    op1 = kfp.components.load_component_from_text(
+      '''
+name: Component name
+implementation:
+  container:
+    image: busybox
+'''
+    )
+
+    @dsl.pipeline()
+    def some_pipeline():
+      op1().set_display_name('Custom name')
+
+    workflow_dict = kfp.compiler.Compiler()._compile(some_pipeline)
+    template = workflow_dict['spec']['templates'][0]
+    self.assertEqual(template['metadata']['annotations']['kubeflow.org/pipelines/task_display_name'], 'Custom name')
+
   def test_op_transformers(self):
     def some_op():
         return dsl.ContainerOp(