docstr: Tests that tasks are generated when upstream is done.
  
  Args:
    use_task_queue: If task queue is enabled, new tasks are only generated if
      a task with the same task_id does not already exist in the queue.
      `use_task_queue=False` is useful to test the case of task generation
      when task queue is empty (for eg: due to orchestrator restart).
    mock_update_external_artifact_type: mock object to the function
      task_gen_utils.update_external_artifact_type
docstr_tok: ["Tests","that","tasks","are","generated","when","upstream","is","done",".","Args",":","use_task_queue",":","If","task","queue","is","enabled",",","new","tasks","are","only","generated","if","a","task","with","the","same","task_id","does","not","already","exist","in","the","queue",".","`","use_task_queue=False","`","is","useful","to","test","the","case","of","task","generation","when","task","queue","is","empty","(","for","eg",":","due","to","orchestrator","restart",")",".","mock_update_external_artifact_type",":","mock","object","to","the","function","task_gen_utils.update_external_artifact_type"]
code: def test_tasks_generated_when_upstream_done(
      self, use_task_queue, mock_update_external_artifact_type
  ):
      """Tests that tasks are generated when upstream is done.
  
      Args:
        use_task_queue: If task queue is enabled, new tasks are only generated if
          a task with the same task_id does not already exist in the queue.
          `use_task_queue=False` is useful to test the case of task generation
          when task queue is empty (for eg: due to orchestrator restart).
        mock_update_external_artifact_type: mock object to the function
          task_gen_utils.update_external_artifact_type
      """
      # Simulate that ExampleGen has already completed successfully.
      test_utils.fake_example_gen_run(
          self._mlmd_connection, self._example_gen, 1, 1
      )
  
      # Generate once. Stats-gen task should be generated.
      [stats_gen_task] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=1,
          num_tasks_generated=1,
          num_new_executions=1,
          num_active_executions=1,
          expected_exec_nodes=[self._stats_gen],
          ignore_update_node_state_tasks=True,
      )
  
      self._mock_service_job_manager.ensure_node_services.assert_called_with(
          mock.ANY, self._example_gen.node_info.id
      )
      self._mock_service_job_manager.reset_mock()
  
      # Finish stats-gen execution.
      self._finish_node_execution(use_task_queue, stats_gen_task)
  
      # Schema-gen should execute next.
      [schema_gen_task] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=2,
          num_tasks_generated=1,
          num_new_executions=1,
          num_active_executions=1,
          expected_exec_nodes=[self._schema_gen],
          ignore_update_node_state_tasks=True,
      )
  
      # Finish schema-gen execution.
      self._finish_node_execution(use_task_queue, schema_gen_task)
  
      # Transform and ExampleValidator should both execute next.
      [
          example_validator_task,
          transform_task,
      ] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=3,
          num_tasks_generated=2,
          num_new_executions=2,
          num_active_executions=2,
          expected_exec_nodes=[
              self._example_validator,
              self._transform,
          ],
          ignore_update_node_state_tasks=True,
      )
  
      # Transform is a "mixed service node".
      self._mock_service_job_manager.ensure_node_services.assert_called_once_with(
          mock.ANY, self._transform.node_info.id
      )
      self._mock_service_job_manager.reset_mock()
  
      # Finish example-validator execution.
      self._finish_node_execution(
          use_task_queue, example_validator_task
      )
  
      # Since transform hasn't finished, trainer will not be triggered yet.
      tasks = self._generate_and_test(
          use_task_queue,
          num_initial_executions=5,
          num_tasks_generated=0 if use_task_queue else 1,
          num_new_executions=0,
          num_active_executions=1,
          expected_exec_nodes=[]
          if use_task_queue
          else [self._transform],
          ignore_update_node_state_tasks=True,
      )
      if not use_task_queue:
          transform_task = tasks[0]
  
      # Finish transform execution.
      self._finish_node_execution(use_task_queue, transform_task)
  
      # Now all trainer upstream nodes are done, so trainer will be triggered.
      [trainer_task] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=5,
          num_tasks_generated=1,
          num_new_executions=1,
          num_active_executions=1,
          expected_exec_nodes=[self._trainer],
          ignore_update_node_state_tasks=True,
      )
  
      # Finish trainer execution.
      self._finish_node_execution(use_task_queue, trainer_task)
  
      # Test task-only dependencies: chore_a and chore_b nodes have no input or
      # output specs but should still be executed in the DAG order.
      [chore_a_task] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=6,
          num_tasks_generated=1,
          num_new_executions=1,
          num_active_executions=1,
          expected_exec_nodes=[self._chore_a],
          ignore_update_node_state_tasks=True,
      )
      self._finish_node_execution(use_task_queue, chore_a_task)
      [chore_b_task] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=7,
          num_tasks_generated=1,
          num_new_executions=1,
          num_active_executions=1,
          expected_exec_nodes=[self._chore_b],
          ignore_update_node_state_tasks=True,
      )
      self._finish_node_execution(use_task_queue, chore_b_task)
  
      # No more components to execute, FinalizePipelineTask should be generated.
      [finalize_task] = self._generate_and_test(
          use_task_queue,
          num_initial_executions=8,
          num_tasks_generated=1,
          num_new_executions=0,
          num_active_executions=0,
          ignore_update_node_state_tasks=True,
      )
      self.assertIsInstance(
          finalize_task, task_lib.FinalizePipelineTask
      )
      self.assertEqual(status_lib.Code.OK, finalize_task.status.code)
      if use_task_queue:
          self.assertTrue(self._task_queue.is_empty())
  
      mock_update_external_artifact_type.assert_called()
  
imports: import itertools
  import os
  from typing import Literal
  import uuid
  from absl.testing import parameterized
  from absl.testing.absltest import mock
  import tensorflow
  from tfx.dsl.compiler import constants
  from tfx.orchestration import data_types_utils
  from tfx.orchestration.experimental.core import constants
  from tfx.orchestration.experimental.core import mlmd_state
  from tfx.orchestration.experimental.core import pipeline_state
  from tfx.orchestration.experimental.core import service_jobs
  from tfx.orchestration.experimental.core import sync_pipeline_task_gen
  from tfx.orchestration.experimental.core import task
  from tfx.orchestration.experimental.core import task_gen_utils
  from tfx.orchestration.experimental.core import task_queue
  from tfx.orchestration.experimental.core import test_utils
  from tfx.orchestration.experimental.core.testing import test_sync_pipeline
  from tfx.orchestration import mlmd_connection_manager
  from tfx.orchestration.portable import runtime_parameter_utils
  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 9 for calling the SyncPipelineTaskGeneratorTest obj's underlying member method code with example usage: obj.test_tasks_generated_when_upstream_done(use_task_queue, mock_update_external_artifact_type) without return types