code: /*
   * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
   */
  
  package io.airbyte.workers.run;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.mockito.ArgumentMatchers.eq;
  import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.verify;
  import static org.mockito.Mockito.when;
  
  import com.google.common.collect.ImmutableMap;
  import io.airbyte.commons.features.FeatureFlags;
  import io.airbyte.commons.json.Jsons;
  import io.airbyte.config.JobConfig.ConfigType;
  import io.airbyte.config.JobResetConnectionConfig;
  import io.airbyte.config.JobSyncConfig;
  import io.airbyte.config.StandardSyncOperation;
  import io.airbyte.config.StandardSyncOutput;
  import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
  import io.airbyte.scheduler.models.Job;
  import io.airbyte.workers.WorkerConstants;
  import io.airbyte.workers.temporal.TemporalClient;
  import io.airbyte.workers.temporal.TemporalResponse;
  import java.io.IOException;
  import java.nio.file.Files;
  import java.nio.file.Path;
  import java.util.List;
  import java.util.UUID;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.Test;
  import org.mockito.ArgumentCaptor;
  
  class TemporalWorkerRunFactoryTest {
  
    private static final UUID CONNECTION_ID = UUID.randomUUID();
    private static final long JOB_ID = 10L;
    private static final int ATTEMPT_ID = 20;
  
    private Path jobRoot;
    private TemporalClient temporalClient;
    private TemporalWorkerRunFactory workerRunFactory;
    private Job job;
  
    @BeforeEach
    void setup() throws IOException {
      final Path workspaceRoot = Files.createTempDirectory(Path.of("/tmp"), "temporal_worker_run_test");
      jobRoot = workspaceRoot.resolve(String.valueOf(JOB_ID)).resolve(String.valueOf(ATTEMPT_ID));
      temporalClient = mock(TemporalClient.class);
      workerRunFactory = new TemporalWorkerRunFactory(
          temporalClient,
          workspaceRoot,
          "unknown airbyte version",
          mock(FeatureFlags.class));
      job = mock(Job.class, RETURNS_DEEP_STUBS);
      when(job.getId()).thenReturn(JOB_ID);
      when(job.getAttemptsCount()).thenReturn(ATTEMPT_ID);
      when(job.getScope()).thenReturn(CONNECTION_ID.toString());
    }
  
    @SuppressWarnings("unchecked")
    @Test
    void testSync() throws Exception {
      when(job.getConfigType()).thenReturn(ConfigType.SYNC);
      final TemporalResponse<StandardSyncOutput> mockResponse = mock(TemporalResponse.class);
      when(temporalClient.submitSync(JOB_ID, ATTEMPT_ID, job.getConfig().getSync(),
          CONNECTION_ID)).thenReturn(mockResponse);
      final WorkerRun workerRun = workerRunFactory.create(job);
      workerRun.call();
      verify(temporalClient).submitSync(JOB_ID, ATTEMPT_ID, job.getConfig().getSync(), CONNECTION_ID);
      assertEquals(jobRoot, workerRun.getJobRoot());
    }
  
    @SuppressWarnings("unchecked")
    @Test
    void testResetConnection() throws Exception {
      final JobResetConnectionConfig resetConfig = new JobResetConnectionConfig()
          .withDestinationDockerImage("airbyte/fusion_reactor")
          .withDestinationConfiguration(Jsons.jsonNode(ImmutableMap.of("a", 1)))
          .withOperationSequence(List.of(new StandardSyncOperation().withName("b")))
          .withConfiguredAirbyteCatalog(new ConfiguredAirbyteCatalog());
      final JobSyncConfig syncConfig = new JobSyncConfig()
          .withSourceDockerImage(WorkerConstants.RESET_JOB_SOURCE_DOCKER_IMAGE_STUB)
          .withDestinationDockerImage(resetConfig.getDestinationDockerImage())
          .withDestinationConfiguration(resetConfig.getDestinationConfiguration())
          .withOperationSequence(List.of(new StandardSyncOperation().withName("b")))
          .withSourceConfiguration(Jsons.emptyObject())
          .withConfiguredAirbyteCatalog(resetConfig.getConfiguredAirbyteCatalog());
      when(job.getConfigType()).thenReturn(ConfigType.RESET_CONNECTION);
      when(job.getConfig().getResetConnection()).thenReturn(resetConfig);
      final TemporalResponse<StandardSyncOutput> mockResponse = mock(TemporalResponse.class);
      when(temporalClient.submitSync(JOB_ID, ATTEMPT_ID, syncConfig, CONNECTION_ID)).thenReturn(mockResponse);
  
      final WorkerRun workerRun = workerRunFactory.create(job);
      workerRun.call();
  
      final ArgumentCaptor<JobSyncConfig> argument = ArgumentCaptor.forClass(JobSyncConfig.class);
      verify(temporalClient).submitSync(eq(JOB_ID), eq(ATTEMPT_ID), argument.capture(), eq(CONNECTION_ID));
      assertEquals(syncConfig, argument.getValue());
      assertEquals(jobRoot, workerRun.getJobRoot());
    }
  
  }
  