docstr: Sanity-checks the construction and dependencies for a 2-step pipeline.
docstr_tok: ["Sanity-checks","the","construction","and","dependencies","for","a","2-step","pipeline","."]
code: def testTwoStepPipeline(self):
      """Sanity-checks the construction and dependencies for a 2-step pipeline."""
      kubeflow_dag_runner.KubeflowDagRunner().run(_two_step_pipeline())
      file_path = os.path.join(self.tmp_dir, "two_step_pipeline.tar.gz")
      self.assertTrue(fileio.exists(file_path))
  
      with tarfile.TarFile.open(file_path).extractfile(
          "pipeline.yaml"
      ) as pipeline_file:
          self.assertIsNotNone(pipeline_file)
          pipeline = yaml.safe_load(pipeline_file)
  
          containers = [
              c
              for c in pipeline["spec"]["templates"]
              if "container" in c
          ]
          self.assertEqual(2, len(containers))
  
          big_query_container = [
              c for c in containers if c["name"] == "bigqueryexamplegen"
          ]
          self.assertEqual(1, len(big_query_container))
          self.assertEqual(
              [
                  "python",
                  "-m",
                  "tfx.orchestration.kubeflow.container_entrypoint",
              ],
              big_query_container[0]["container"]["command"],
          )
          self.assertIn(
              "--tfx_ir", big_query_container[0]["container"]["args"]
          )
          self.assertIn(
              "--node_id", big_query_container[0]["container"]["args"]
          )
          self._compare_tfx_ir_against_testdata(
              big_query_container[0]["container"]["args"],
              "two_step_pipeline_post_dehydrate_ir.json",
          )
  
          statistics_gen_container = [
              c for c in containers if c["name"] == "statisticsgen"
          ]
          self.assertEqual(1, len(statistics_gen_container))
  
          # Ensure the pod labels are correctly appended.
          metadata = [
              c["metadata"]
              for c in pipeline["spec"]["templates"]
              if "dag" not in c
          ]
          for m in metadata:
              self.assertEqual(
                  "tfx", m["labels"][telemetry_utils.LABEL_KFP_SDK_ENV]
              )
  
          # Ensure dependencies between components are captured.
          dag = [c for c in pipeline["spec"]["templates"] if "dag" in c]
          self.assertEqual(1, len(dag))
  
          self.assertEqual(
              {
                  "tasks": [
                      {
                          "name": "bigqueryexamplegen",
                          "template": "bigqueryexamplegen",
                          "arguments": {
                              "parameters": [
                                  {
                                      "name": "input_config",
                                      "value": "{{inputs.parameters.input_config}}",
                                  },
                                  {
                                      "name": "pipeline-root",
                                      "value": "{{inputs.parameters.pipeline-root}}",
                                  },
                              ]
                          },
                      },
                      {
                          "name": "statisticsgen",
                          "template": "statisticsgen",
                          "arguments": {
                              "parameters": [
                                  {
                                      "name": "pipeline-root",
                                      "value": "{{inputs.parameters.pipeline-root}}",
                                  }
                              ]
                          },
                          "dependencies": ["bigqueryexamplegen"],
                      },
                  ]
              },
              dag[0]["dag"],
          )
  
imports: import json
  import os
  import tarfile
  from typing import List
  from kfp import onprem
  import tensorflow
  from tfx.components.statistics_gen import component
  from tfx.dsl.component.experimental import executor_specs
  from tfx.dsl.component.experimental.annotations import Parameter
  from tfx.dsl.component.experimental.decorators import component
  from tfx.dsl.components.base import base_component
  from tfx.dsl.io import fileio
  from tfx.extensions.google_cloud_big_query.example_gen import component
  from tfx.orchestration import data_types
  from tfx.orchestration import pipeline
  from tfx.orchestration.kubeflow import kubeflow_dag_runner
  from tfx.orchestration.kubeflow.decorators import FinalStatusStr
  from tfx.proto import example_gen_pb2
  from tfx.types import component_spec
  from tfx.utils import telemetry_utils
  from tfx.utils import test_case_utils
  import yaml
  from ml_metadata.proto import metadata_store_pb2
label_desc: Use image node_id 3 for calling the KubeflowDagRunnerTest obj's underlying member method code with example usage: obj.testTwoStepPipeline() without return types