chore(sdk): move pipeline test samples closer to compiler unit tests. (#7849)

* chore(sdk): move pipeline test samples closer to compiler unit tests.

* explicitly list test files

* remove dead code

sdk/python/kfp/cli/cli_test.py

@@ -14,17 +14,15 @@
 
 import functools
 import itertools
-import json
 import os
 import re
 import subprocess
 import tempfile
 import unittest
-from typing import Any, Dict, List, Optional
+from typing import List
 from unittest import mock
 
 import click
-import yaml
 from absl.testing import parameterized
 from click import testing
 from kfp.cli import cli
@@ -53,37 +51,6 @@ class TestCliNounAliases(unittest.TestCase):
                          result.output)
 
 
-def _ignore_kfp_version_helper(spec: Dict[str, Any]) -> Dict[str, Any]:
-    """Ignores kfp sdk versioning in command.
-
-    Takes in a YAML input and ignores the kfp sdk versioning in command
-    for comparison between compiled file and goldens.
-    """
-    pipeline_spec = spec.get('pipelineSpec', spec)
-
-    if 'executors' in pipeline_spec['deploymentSpec']:
-        for executor in pipeline_spec['deploymentSpec']['executors']:
-            pipeline_spec['deploymentSpec']['executors'][
-                executor] = yaml.safe_load(
-                    re.sub(
-                        r"'kfp==(\d+).(\d+).(\d+)(-[a-z]+.\d+)?'", 'kfp',
-                        yaml.dump(
-                            pipeline_spec['deploymentSpec']['executors']
-                            [executor],
-                            sort_keys=True)))
-    return spec
-
-
-def load_compiled_file(filename: str) -> Dict[str, Any]:
-    with open(filename, 'r') as f:
-        contents = yaml.safe_load(f)
-        pipeline_spec = contents[
-            'pipelineSpec'] if 'pipelineSpec' in contents else contents
-        # ignore the sdkVersion
-        del pipeline_spec['sdkVersion']
-        return _ignore_kfp_version_helper(contents)
-
-
 class TestAliases(unittest.TestCase):
 
     @classmethod
@@ -175,19 +142,6 @@ class TestCliVersion(unittest.TestCase):
         self.assertTrue(matches)
 
 
-COMPILER_CLI_TEST_DATA_DIR = os.path.join(
-    os.path.dirname(os.path.dirname(__file__)), 'compiler_cli_tests',
-    'test_data')
-
-SPECIAL_TEST_PY_FILES = {'two_step_pipeline.py'}
-
-TEST_PY_FILES = {
-    file.split('.')[0]
-    for file in os.listdir(COMPILER_CLI_TEST_DATA_DIR)
-    if ".py" in file and file not in SPECIAL_TEST_PY_FILES
-}
-
-
 class TestDslCompile(parameterized.TestCase):
 
     def invoke(self, args: List[str]) -> testing.Result:
@@ -205,54 +159,6 @@ class TestDslCompile(parameterized.TestCase):
             catch_exceptions=False,
             obj={})
 
-    def _test_compile_py_to_yaml(
-            self,
-            file_base_name: str,
-            additional_arguments: Optional[List[str]] = None) -> None:
-        py_file = os.path.join(COMPILER_CLI_TEST_DATA_DIR,
-                               f'{file_base_name}.py')
-
-        golden_compiled_file = os.path.join(COMPILER_CLI_TEST_DATA_DIR,
-                                            f'{file_base_name}.yaml')
-
-        if additional_arguments is None:
-            additional_arguments = []
-
-        with tempfile.TemporaryDirectory() as tmpdir:
-            generated_compiled_file = os.path.join(
-                tmpdir, f'{file_base_name}-pipeline.yaml')
-
-            result = self.invoke(
-                ['--py', py_file, '--output', generated_compiled_file] +
-                additional_arguments)
-
-            self.assertEqual(result.exit_code, 0)
-
-            compiled = load_compiled_file(generated_compiled_file)
-
-        golden = load_compiled_file(golden_compiled_file)
-        self.assertEqual(golden, compiled)
-
-    def test_two_step_pipeline(self):
-        self._test_compile_py_to_yaml(
-            'two_step_pipeline',
-            ['--pipeline-parameters', '{"text":"Hello KFP!"}'])
-
-    def test_two_step_pipeline_failure_parameter_parse(self):
-        with self.assertRaisesRegex(json.decoder.JSONDecodeError,
-                                    r"Unterminated string starting at:"):
-            self._test_compile_py_to_yaml(
-                'two_step_pipeline',
-                ['--pipeline-parameters', '{"text":"Hello KFP!}'])
-
-    @parameterized.parameters(TEST_PY_FILES)
-    def test_compile_pipelines(self, file: str):
-
-        # To update all golden snapshots:
-        # for f in test_data/*.py ; do python3 "$f" ; done
-
-        self._test_compile_py_to_yaml(file)
-
     def test_deprecated_command_is_found(self):
         result = self.invoke_deprecated(['--help'])
         self.assertEqual(result.exit_code, 0)

sdk/python/kfp/compiler/compiler_test.py

@@ -14,18 +14,22 @@
 
 import json
 import os
+import re
 import tempfile
 import unittest
+from typing import Any, Dict, List, Optional
 
+import yaml
 from absl.testing import parameterized
+from click import testing
 from google.protobuf import json_format
 from kfp import components
 from kfp import dsl
+from kfp.cli import cli
 from kfp.compiler import compiler
 from kfp.components.types import type_utils
 from kfp.dsl import PipelineTaskFinalStatus
 from kfp.pipeline_spec import pipeline_spec_pb2
-import yaml
 
 VALID_PRODUCER_COMPONENT_SAMPLE = components.load_component_from_text("""
     name: producer
@@ -862,5 +866,116 @@ class TestWriteIrToFile(unittest.TestCase):
                                                  temp_filepath)
 
 
+SAMPLE_PIPELINES_TEST_DATA_DIR = os.path.join(
+    os.path.dirname(__file__), 'test_data')
+
+SAMPLE_TEST_PY_FILES = [
+    'pipeline_with_importer',
+    'pipeline_with_ontology',
+    'pipeline_with_if_placeholder',
+    'pipeline_with_concat_placeholder',
+    'pipeline_with_resource_spec',
+    'pipeline_with_various_io_types',
+    'pipeline_with_reused_component',
+    'pipeline_with_after',
+    'pipeline_with_condition',
+    'pipeline_with_nested_conditions',
+    'pipeline_with_nested_conditions_yaml',
+    'pipeline_with_loops',
+    'pipeline_with_nested_loops',
+    'pipeline_with_loops_and_conditions',
+    'pipeline_with_params_containing_format',
+    'lightweight_python_functions_v2_pipeline',
+    'lightweight_python_functions_v2_with_outputs',
+    'xgboost_sample_pipeline',
+    'pipeline_with_metrics_outputs',
+    'pipeline_with_exit_handler',
+    'pipeline_with_env',
+    'v2_component_with_optional_inputs',
+    'pipeline_with_gcpc_types',
+    'pipeline_with_placeholders',
+    'pipeline_with_task_final_status',
+    'pipeline_with_task_final_status_yaml',
+    'v2_component_with_pip_index_urls',
+]
+
+
+class TestDslCompileSamplePipelines(parameterized.TestCase):
+
+    def _invoke(self, args: List[str]) -> testing.Result:
+        starting_args = ['dsl', 'compile']
+        args = starting_args + args
+        runner = testing.CliRunner()
+        return runner.invoke(
+            cli=cli.cli, args=args, catch_exceptions=False, obj={})
+
+    def _ignore_kfp_version_helper(self, spec: Dict[str,
+                                                    Any]) -> Dict[str, Any]:
+        """Ignores kfp sdk versioning in command.
+
+        Takes in a YAML input and ignores the kfp sdk versioning in
+        command for comparison between compiled file and goldens.
+        """
+        pipeline_spec = spec.get('pipelineSpec', spec)
+
+        if 'executors' in pipeline_spec['deploymentSpec']:
+            for executor in pipeline_spec['deploymentSpec']['executors']:
+                pipeline_spec['deploymentSpec']['executors'][
+                    executor] = yaml.safe_load(
+                        re.sub(
+                            r"'kfp==(\d+).(\d+).(\d+)(-[a-z]+.\d+)?'", 'kfp',
+                            yaml.dump(
+                                pipeline_spec['deploymentSpec']['executors']
+                                [executor],
+                                sort_keys=True)))
+        return spec
+
+    def _load_compiled_file(self, filename: str) -> Dict[str, Any]:
+        with open(filename, 'r') as f:
+            contents = yaml.safe_load(f)
+            pipeline_spec = contents[
+                'pipelineSpec'] if 'pipelineSpec' in contents else contents
+            # ignore the sdkVersion
+            del pipeline_spec['sdkVersion']
+            return self._ignore_kfp_version_helper(contents)
+
+    def _test_compile_py_to_yaml(
+            self,
+            file_base_name: str,
+            additional_arguments: Optional[List[str]] = None) -> None:
+        py_file = os.path.join(SAMPLE_PIPELINES_TEST_DATA_DIR,
+                               f'{file_base_name}.py')
+
+        golden_compiled_file = os.path.join(SAMPLE_PIPELINES_TEST_DATA_DIR,
+                                            f'{file_base_name}.yaml')
+
+        if additional_arguments is None:
+            additional_arguments = []
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            generated_compiled_file = os.path.join(
+                tmpdir, f'{file_base_name}-pipeline.yaml')
+
+            result = self._invoke(
+                ['--py', py_file, '--output', generated_compiled_file] +
+                additional_arguments)
+
+            self.assertEqual(result.exit_code, 0)
+
+            compiled = self._load_compiled_file(generated_compiled_file)
+
+        golden = self._load_compiled_file(golden_compiled_file)
+        self.assertEqual(golden, compiled)
+
+    def test_two_step_pipeline(self):
+        self._test_compile_py_to_yaml(
+            'two_step_pipeline',
+            ['--pipeline-parameters', '{"text":"Hello KFP!"}'])
+
+    @parameterized.parameters(SAMPLE_TEST_PY_FILES)
+    def test_compile_pipelines(self, file: str):
+        self._test_compile_py_to_yaml(file)
+
+
 if __name__ == '__main__':
     unittest.main()

sdk/python/kfp/compiler/test_data/README.md

@@ -0,0 +1,7 @@
+# Pipeline samples for compiler unit tests.
+
+To update all golden snapshots:
+
+```bash
+for f in test_data/*.py ; do python3 "$f" ; done
+```