fix tests

kubeflow · Aug 31, 2022 · 851a50d · 851a50d
1 parent 7e6c52e
commit 851a50d
Showing 1 changed file with 4 additions and 9 deletions.
diff --git a/sdk/python/kfp/components/importer_node_test.py b/sdk/python/kfp/components/importer_node_test.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 import unittest
 
-from kfp import compiler
 from kfp import dsl
 from kfp.components import importer_node
 from kfp.components.types.artifact_types import Dataset
@@ -35,8 +34,7 @@ def my_pipeline(meta_inp: str):
                     'string-2': meta_inp
                 })
 
-        pipeline_spec = compiler.Compiler()._create_pipeline(
-            pipeline_func=my_pipeline)
+        pipeline_spec = my_pipeline.pipeline_spec
         input_keys = list(pipeline_spec.components['comp-importer']
                           .input_definitions.parameters.keys())
         self.assertIn('meta_inp', input_keys)
@@ -63,8 +61,7 @@ def my_pipeline(meta_inp: str):
                     meta_inp: meta_inp
                 })
 
-        pipeline_spec = compiler.Compiler()._create_pipeline(
-            pipeline_func=my_pipeline)
+        pipeline_spec = my_pipeline.pipeline_spec
         input_keys = list(pipeline_spec.components['comp-importer']
                           .input_definitions.parameters.keys())
         self.assertIn('meta_inp', input_keys)
@@ -102,8 +99,7 @@ def my_pipeline():
                     'string-2': task2.output
                 })
 
-        pipeline_spec = compiler.Compiler()._create_pipeline(
-            pipeline_func=my_pipeline)
+        pipeline_spec = my_pipeline.pipeline_spec
         input_keys = list(pipeline_spec.components['comp-importer']
                           .input_definitions.parameters.keys())
         self.assertIn('Output', input_keys)
@@ -136,8 +132,7 @@ def my_pipeline():
                     task.output: task.output
                 })
 
-        pipeline_spec = compiler.Compiler()._create_pipeline(
-            pipeline_func=my_pipeline)
+        pipeline_spec = my_pipeline.pipeline_spec
         input_keys = list(pipeline_spec.components['comp-importer']
                           .input_definitions.parameters.keys())
         self.assertIn('Output', input_keys)