code: def test_task_handling(
      self, mock_publish, mock_record_state_change_time
  ):
      collector = _Collector()
  
      # Register a fake task scheduler.
      ts.TaskSchedulerRegistry.register(
          self._type_url,
          functools.partial(
              _FakeTaskScheduler,
              block_nodes={"Trainer", "Transform", "Pusher"},
              collector=collector,
          ),
      )
  
      task_queue = tq.TaskQueue()
  
      # Enqueue some tasks.
      trainer_exec_task = _test_exec_node_task(
          "Trainer", "test-pipeline", pipeline=self._pipeline
      )
      task_queue.enqueue(trainer_exec_task)
      task_queue.enqueue(
          _test_cancel_node_task("Trainer", "test-pipeline")
      )
  
      with self._task_manager(task_queue) as task_manager:
          # Enqueue more tasks after task manager starts.
          transform_exec_task = _test_exec_node_task(
              "Transform", "test-pipeline", pipeline=self._pipeline
          )
          task_queue.enqueue(transform_exec_task)
          evaluator_exec_task = _test_exec_node_task(
              "Evaluator", "test-pipeline", pipeline=self._pipeline
          )
          task_queue.enqueue(evaluator_exec_task)
          task_queue.enqueue(
              _test_cancel_node_task("Transform", "test-pipeline")
          )
          pusher_exec_task = _test_exec_node_task(
              "Pusher", "test-pipeline", pipeline=self._pipeline
          )
          task_queue.enqueue(pusher_exec_task)
          task_queue.enqueue(
              _test_cancel_node_task(
                  "Pusher", "test-pipeline", pause=True
              )
          )
  
      self.assertTrue(task_manager.done())
      self.assertIsNone(task_manager.exception())
  
      # Ensure that all exec and cancellation tasks were processed correctly.
      self.assertCountEqual(
          [
              trainer_exec_task,
              transform_exec_task,
              evaluator_exec_task,
              pusher_exec_task,
          ],
          collector.scheduled_tasks,
      )
      self.assertCountEqual(
          [
              trainer_exec_task,
              transform_exec_task,
              pusher_exec_task,
          ],
          collector.cancelled_tasks,
      )
  
      result_ok = ts.TaskSchedulerResult(
          status=status_lib.Status(
              code=status_lib.Code.OK,
              message="_FakeTaskScheduler result",
          )
      )
      result_cancelled = ts.TaskSchedulerResult(
          status=status_lib.Status(
              code=status_lib.Code.CANCELLED,
              message="_FakeTaskScheduler result",
          )
      )
      mock_publish.assert_has_calls(
          [
              mock.call(
                  mlmd_handle=mock.ANY,
                  task=trainer_exec_task,
                  result=result_cancelled,
              ),
              mock.call(
                  mlmd_handle=mock.ANY,
                  task=transform_exec_task,
                  result=result_cancelled,
              ),
              mock.call(
                  mlmd_handle=mock.ANY,
                  task=evaluator_exec_task,
                  result=result_ok,
              ),
          ],
          any_order=True,
      )
      # It is expected that publish is not called for Pusher because it was
      # cancelled with pause=True so there must be only 3 calls.
      self.assertLen(mock_publish.mock_calls, 3)
      self.assertLen(mock_record_state_change_time.mock_calls, 3)
  
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 3 for calling the TaskManagerTest obj's underlying member method code with example usage: obj.test_task_handling(mock_publish, mock_record_state_change_time) without return types