class_docstr: Test end-to-end from task generation to publication of results to MLMD.
class_docstr_tok: ["Test","end-to-end","from","task","generation","to","publication","of","results","to","MLMD","."]
code: def setUp(self):
      super().setUp()
      pipeline_root = os.path.join(
          os.environ.get(
              "TEST_UNDECLARED_OUTPUTS_DIR", self.get_temp_dir()
          ),
          self.id(),
      )
  
      # Makes sure multiple connections within a test always connect to the same
      # MLMD instance.
      metadata_path = os.path.join(
          pipeline_root, "metadata", "metadata.db"
      )
      self._mlmd_connection_manager = (
          mlmd_cm.MLMDConnectionManager.sqlite(metadata_path)
      )
      self.enter_context(self._mlmd_connection_manager)
      self._mlmd_connection = (
          self._mlmd_connection_manager.primary_mlmd_handle
      )
  
      # Sets up the pipeline.
      pipeline = test_async_pipeline.create_pipeline()
  
      # Extracts components.
      self._example_gen = pipeline.nodes[0].pipeline_node
      self._transform = pipeline.nodes[1].pipeline_node
      self._trainer = pipeline.nodes[2].pipeline_node
  
      # Pack deployment config for testing.
      deployment_config = pipeline_pb2.IntermediateDeploymentConfig()
      executor_spec = pipeline_pb2.ExecutorSpec.PythonClassExecutorSpec(
          class_path="fake.ClassPath"
      )
      deployment_config.executor_specs[self._trainer.node_info.id].Pack(
          executor_spec
      )
      deployment_config.executor_specs[
          self._transform.node_info.id
      ].Pack(executor_spec)
      self._type_url = deployment_config.executor_specs[
          self._trainer.node_info.id
      ].type_url
      pipeline.deployment_config.Pack(deployment_config)
      self._pipeline = pipeline
      self._pipeline_info = pipeline.pipeline_info
      self._pipeline_runtime_spec = pipeline.runtime_spec
      self._pipeline_runtime_spec.pipeline_root.field_value.string_value = (
          pipeline_root
      )
  
      ts.TaskSchedulerRegistry.clear()
      self._task_queue = tq.TaskQueue()
  
      # Run fake example-gen to prepare downstreams component triggers.
      test_utils.fake_example_gen_run(
          self._mlmd_connection, self._example_gen, 1, 1
      )
  
      # Task generator should produce three tasks for transform. The first one is
      # UpdateNodeStateTask with state RUNNING, the second one is ExecNodeTask
      # and the third one is UpdateNodeStateTask with state STARTED
      with self._mlmd_connection_manager as mlmd_connection_manager:
          m = mlmd_connection_manager.primary_mlmd_handle
          pipeline_state = pstate.PipelineState.new(m, self._pipeline)
          tasks = asptg.AsyncPipelineTaskGenerator(
              mlmd_connection_manager,
              self._task_queue.contains_task_id,
              service_jobs.DummyServiceJobManager(),
          ).generate(pipeline_state)
      self.assertLen(tasks, 3)
      self.assertIsInstance(tasks[0], task_lib.UpdateNodeStateTask)
      self.assertEqual("my_transform", tasks[0].node_uid.node_id)
      self.assertEqual(pstate.NodeState.RUNNING, tasks[0].state)
      self.assertIsInstance(tasks[1], task_lib.ExecNodeTask)
      self.assertEqual("my_transform", tasks[1].node_uid.node_id)
      self.assertTrue(os.path.exists(tasks[1].stateful_working_dir))
      self.assertTrue(os.path.exists(tasks[1].tmp_dir))
      self.assertIsInstance(tasks[2], task_lib.UpdateNodeStateTask)
      self.assertEqual("my_trainer", tasks[2].node_uid.node_id)
      self.assertEqual(pstate.NodeState.STARTED, tasks[2].state)
  
      self._task = tasks[1]
      self._output_artifact_uri = self._task.output_artifacts[
          "transform_graph"
      ][0].uri
      self.assertTrue(os.path.exists(self._output_artifact_uri))
      self._task_queue.enqueue(self._task)
  
      # There should be 1 active execution in MLMD.
      with self._mlmd_connection as m:
          executions = m.store.get_executions()
      active_executions = [
          e
          for e in executions
          if e.last_known_state == metadata_store_pb2.Execution.RUNNING
      ]
      self.assertLen(active_executions, 1)
  
      # Active execution id.
      self._execution_id = active_executions[0].id
  
imports: import contextlib
  import functools
  import os
  import threading
  import time
  from absl import logging
  from absl.testing.absltest import mock
  import tensorflow
  from tfx.orchestration import data_types_utils
  from tfx.orchestration import metadata
  from tfx.orchestration.experimental.core import async_pipeline_task_gen
  from tfx.orchestration.experimental.core import constants
  from tfx.orchestration.experimental.core import pipeline_state
  from tfx.orchestration.experimental.core import post_execution_utils
  from tfx.orchestration.experimental.core import service_jobs
  from tfx.orchestration.experimental.core import task
  from tfx.orchestration.experimental.core import task_manager
  from tfx.orchestration.experimental.core import task_queue
  from tfx.orchestration.experimental.core import task_scheduler
  from tfx.orchestration.experimental.core import test_utils
  from tfx.orchestration.experimental.core.testing import test_async_pipeline
  from tfx.orchestration import mlmd_connection_manager
  from tfx.proto.orchestration import execution_result_pb2
  from tfx.proto.orchestration import pipeline_pb2
  from tfx.utils import status
  from ml_metadata.proto import metadata_store_pb2
label_desc: Use image node_id 1 for calling the TaskManagerE2ETest obj's underlying member method code with example usage: obj.setUp() without return types