content: /**
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  package org.apache.atlas.repository.store.graph.v2;
  
  import org.apache.atlas.AtlasErrorCode;
  import org.apache.atlas.exception.AtlasBaseException;
  import org.apache.atlas.kafka.NotificationProvider;
  import org.apache.atlas.model.impexp.AtlasAsyncImportRequest;
  import org.apache.atlas.model.impexp.AtlasImportResult;
  import org.apache.atlas.model.instance.AtlasEntity;
  import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
  import org.apache.atlas.model.notification.MessageSource;
  import org.apache.atlas.notification.NotificationException;
  import org.apache.atlas.notification.NotificationInterface;
  import org.apache.atlas.repository.impexp.AsyncImportService;
  import org.apache.atlas.repository.store.graph.v2.asyncimport.ImportTaskListener;
  import org.janusgraph.diskstorage.locking.PermanentLockingException;
  import org.mockito.Mock;
  import org.mockito.MockedStatic;
  import org.mockito.Mockito;
  import org.mockito.MockitoAnnotations;
  import org.testng.annotations.BeforeMethod;
  import org.testng.annotations.DataProvider;
  import org.testng.annotations.Test;
  
  import java.util.Collections;
  
  import static org.mockito.ArgumentMatchers.any;
  import static org.mockito.Mockito.anyList;
  import static org.mockito.Mockito.anyString;
  import static org.mockito.Mockito.doNothing;
  import static org.mockito.Mockito.doThrow;
  import static org.mockito.Mockito.eq;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.never;
  import static org.mockito.Mockito.spy;
  import static org.mockito.Mockito.times;
  import static org.mockito.Mockito.verify;
  import static org.mockito.Mockito.when;
  import static org.testng.Assert.assertEquals;
  import static org.testng.Assert.assertNotNull;
  import static org.testng.Assert.assertSame;
  import static org.testng.Assert.expectThrows;
  
  public class AsyncImportTaskExecutorTest {
      @Mock
      private AsyncImportService importService;
  
      @Mock
      private NotificationInterface notificationInterface;
  
      @Mock
      private ImportTaskListener importTaskListener;
  
      @Mock
      private MessageSource messageSource;
  
      private AsyncImportTaskExecutor asyncImportTaskExecutor;
  
      @BeforeMethod
      public void setup() {
          MockitoAnnotations.openMocks(this);
  
          try (MockedStatic<NotificationProvider> mockedStatic = Mockito.mockStatic(NotificationProvider.class)) {
              mockedStatic.when(NotificationProvider::get).thenReturn(notificationInterface);
  
              when(messageSource.getSource()).thenReturn("AsyncImportTaskPublisher");
  
              asyncImportTaskExecutor = new AsyncImportTaskExecutor(importService, importTaskListener);
          }
      }
  
      @Test
      void testRunSuccess() throws AtlasBaseException {
          AtlasImportResult  mockResult             = mock(AtlasImportResult.class);
          EntityImportStream mockEntityImportStream = mock(EntityImportStream.class);
  
          when(mockEntityImportStream.getMd5Hash()).thenReturn("import-md5-hash");
          when(mockEntityImportStream.size()).thenReturn(5);
          when(mockEntityImportStream.getCreationOrder()).thenReturn(Collections.emptyList());
          when(mockEntityImportStream.hasNext()).thenReturn(false);
  
          when(importService.fetchImportRequestByImportId("import-md5-hash")).thenReturn(null);
          doNothing().when(importService).saveImportRequest(any(AtlasAsyncImportRequest.class));
  
          AtlasAsyncImportRequest result = asyncImportTaskExecutor.run(mockResult, mockEntityImportStream);
  
          assertNotNull(result);
          assertSame(result.getStatus(), AtlasAsyncImportRequest.ImportStatus.STAGING);
          verify(mockEntityImportStream).close();
          verify(importService).saveImportRequest(any(AtlasAsyncImportRequest.class));
      }
  
      @Test
      void testRunDuplicateRequestInWaitingStatus() throws AtlasBaseException {
          AtlasImportResult  mockResult             = mock(AtlasImportResult.class);
          EntityImportStream mockEntityImportStream = mock(EntityImportStream.class);
  
          when(mockEntityImportStream.getMd5Hash()).thenReturn("import-md5");
          when(mockEntityImportStream.size()).thenReturn(10);
          when(mockEntityImportStream.getCreationOrder()).thenReturn(Collections.emptyList());
  
          AtlasAsyncImportRequest mockRequest = mock(AtlasAsyncImportRequest.class);
          when(mockRequest.getStatus()).thenReturn(AtlasAsyncImportRequest.ImportStatus.WAITING);
          when(importService.fetchImportRequestByImportId("import-md5")).thenReturn(mockRequest);
  
          doNothing().when(mockEntityImportStream).close();
  
          AtlasAsyncImportRequest result = asyncImportTaskExecutor.run(mockResult, mockEntityImportStream);
  
          assertNotNull(result);
          assertSame(result.getStatus(), AtlasAsyncImportRequest.ImportStatus.WAITING);
          verify(mockEntityImportStream).close();
          verify(importService, never()).saveImportRequest(any());
          verify(importService, never()).updateImportRequest(any());
  
          // Verify that skipToPosition and publishImportRequest are NOT called
          AsyncImportTaskExecutor spyPublisher = spy(asyncImportTaskExecutor);
  
          doNothing().when(spyPublisher).skipToStartEntityPosition(any(), any());
          doNothing().when(spyPublisher).publishImportRequest(any(), any());
  
          verify(spyPublisher, never()).skipToStartEntityPosition(mockRequest, mockEntityImportStream);
          verify(spyPublisher, never()).publishImportRequest(mockRequest, mockEntityImportStream);
      }
  
      @Test
      void testRunDuplicateRequestInProcessingStatus() throws AtlasBaseException {
          AtlasImportResult  mockResult             = mock(AtlasImportResult.class);
          EntityImportStream mockEntityImportStream = mock(EntityImportStream.class);
  
          when(mockEntityImportStream.getMd5Hash()).thenReturn("import-md5");
          when(mockEntityImportStream.size()).thenReturn(10);
          when(mockEntityImportStream.getCreationOrder()).thenReturn(Collections.emptyList());
  
          AtlasAsyncImportRequest mockRequest = mock(AtlasAsyncImportRequest.class);
  
          when(mockRequest.getStatus()).thenReturn(AtlasAsyncImportRequest.ImportStatus.PROCESSING);
          when(importService.fetchImportRequestByImportId("import-md5")).thenReturn(mockRequest);
  
          doNothing().when(mockEntityImportStream).close();
  
          AtlasAsyncImportRequest result = asyncImportTaskExecutor.run(mockResult, mockEntityImportStream);
  
          assertNotNull(result);
          assertSame(result.getStatus(), AtlasAsyncImportRequest.ImportStatus.PROCESSING);
          verify(mockEntityImportStream).close();
          verify(importService, never()).saveImportRequest(any());
          verify(importService, never()).updateImportRequest(any());
  
          // Verify that skipToPosition and publishImportRequest are NOT called
          AsyncImportTaskExecutor spyPublisher = spy(asyncImportTaskExecutor);
  
          doNothing().when(spyPublisher).skipToStartEntityPosition(any(), any());
          doNothing().when(spyPublisher).publishImportRequest(any(), any());
  
          verify(spyPublisher, never()).skipToStartEntityPosition(mockRequest, mockEntityImportStream);
          verify(spyPublisher, never()).publishImportRequest(mockRequest, mockEntityImportStream);
      }
  
      @Test
      void testPublishImportRequestHappyPath() throws AtlasBaseException {
          AtlasAsyncImportRequest mockImportRequest      = mock(AtlasAsyncImportRequest.class);
          AtlasImportResult       mockResult             = mock(AtlasImportResult.class);
          EntityImportStream      mockEntityImportStream = mock(EntityImportStream.class);
  
          when(mockImportRequest.getTopicName()).thenReturn("test-topic");
          when(mockImportRequest.getImportId()).thenReturn("import-md5");
          when(mockImportRequest.getImportDetails()).thenReturn(new AtlasAsyncImportRequest.ImportDetails());
          when(mockImportRequest.getImportResult()).thenReturn(mockResult);
          when(mockResult.getUserName()).thenReturn("test-user-1");
          when(mockEntityImportStream.getTypesDef()).thenReturn(null);
  
          asyncImportTaskExecutor.publishImportRequest(mockImportRequest, mockEntityImportStream);
  
          verify(importService).updateImportRequest(mockImportRequest);
          verify(notificationInterface).closeProducer(NotificationInterface.NotificationType.ASYNC_IMPORT, "test-topic");
          verify(importTaskListener).onReceiveImportRequest(mockImportRequest);
      }
  
      @Test
      void testPublishImportRequestTypeDefNotificationException() throws AtlasBaseException, NotificationException {
          AtlasAsyncImportRequest mockImportRequest      = mock(AtlasAsyncImportRequest.class);
          AtlasImportResult       mockResult             = mock(AtlasImportResult.class);
          EntityImportStream      mockEntityImportStream = mock(EntityImportStream.class);
  
          when(mockImportRequest.getTopicName()).thenReturn("test-topic");
          when(mockImportRequest.getImportId()).thenReturn("import-md5");
          when(mockImportRequest.getImportResult()).thenReturn(mockResult);
          when(mockResult.getUserName()).thenReturn("test-user-1");
          when(mockEntityImportStream.getTypesDef()).thenReturn(null);
  
          doThrow(new NotificationException(new Exception("some notification exception")))
                  .when(notificationInterface)
                  .send(eq("test-topic"), anyList(), any());
  
          try {
              asyncImportTaskExecutor.publishImportRequest(mockImportRequest, mockEntityImportStream);
          } catch (AtlasBaseException ignored) {
              // Ignored for this test
          }
  
          verify(importTaskListener, never()).onReceiveImportRequest(any(AtlasAsyncImportRequest.class));
          verify(notificationInterface).closeProducer(NotificationInterface.NotificationType.ASYNC_IMPORT, "test-topic");
      }
  
      @Test
      void testPublishEntityNotificationHappyPath() throws NotificationException {
          AtlasAsyncImportRequest mockImportRequest      = mock(AtlasAsyncImportRequest.class);
          AtlasImportResult       mockResult             = mock(AtlasImportResult.class);
          EntityImportStream      mockEntityImportStream = mock(EntityImportStream.class);
          AtlasEntityWithExtInfo  mockEntityWithExtInfo  = mock(AtlasEntityWithExtInfo.class);
  
          when(mockImportRequest.getTopicName()).thenReturn("test-topic");
          when(mockImportRequest.getImportId()).thenReturn("import-id");
          when(mockImportRequest.getImportDetails()).thenReturn(new AtlasAsyncImportRequest.ImportDetails());
          when(mockImportRequest.getImportTrackingInfo()).thenReturn(new AtlasAsyncImportRequest.ImportTrackingInfo());
          when(mockImportRequest.getImportResult()).thenReturn(mockResult);
          when(mockResult.getUserName()).thenReturn("test-user-1");
          when(mockEntityImportStream.hasNext()).thenReturn(true, false); // One entity in the stream
          when(mockEntityImportStream.getNextEntityWithExtInfo()).thenReturn(mockEntityWithExtInfo);
          when(mockEntityImportStream.getPosition()).thenReturn(1);
  
          AtlasEntity mockEntity = mock(AtlasEntity.class);
  
          when(mockEntityWithExtInfo.getEntity()).thenReturn(mockEntity);
          when(mockEntity.getGuid()).thenReturn("entity-guid");
  
          asyncImportTaskExecutor.publishEntityNotification(mockImportRequest, mockEntityImportStream);
  
          verify(notificationInterface).send(eq("test-topic"), anyList(), any());
          verify(mockEntityImportStream).onImportComplete("entity-guid");
          verify(importService).updateImportRequest(mockImportRequest);
          assertEquals(mockImportRequest.getImportTrackingInfo().getStartEntityPosition(), 1);
          assertEquals(mockImportRequest.getImportDetails().getPublishedEntityCount(), 1);
      }
  
      @Test
      void testPublishEntityNotificationNullEntity() throws NotificationException {
          AtlasAsyncImportRequest mockImportRequest      = mock(AtlasAsyncImportRequest.class);
          EntityImportStream      mockEntityImportStream = mock(EntityImportStream.class);
  
          when(mockImportRequest.getImportDetails()).thenReturn(new AtlasAsyncImportRequest.ImportDetails());
          when(mockImportRequest.getImportTrackingInfo()).thenReturn(new AtlasAsyncImportRequest.ImportTrackingInfo());
          when(mockEntityImportStream.hasNext()).thenReturn(true, false); // One entity in the stream
          when(mockEntityImportStream.getNextEntityWithExtInfo()).thenReturn(null);
          when(mockEntityImportStream.getPosition()).thenReturn(1);
  
          asyncImportTaskExecutor.publishEntityNotification(mockImportRequest, mockEntityImportStream);
  
          verify(notificationInterface, never()).send(anyString(), anyList(), any());
          verify(mockEntityImportStream, never()).onImportComplete(anyString());
          verify(importService).updateImportRequest(mockImportRequest);
          assertEquals(mockImportRequest.getImportTrackingInfo().getStartEntityPosition(), 1);
          assertEquals(mockImportRequest.getImportDetails().getPublishedEntityCount(), 0);
      }
  
      @Test
      void testPublishEntityNotificationExceptionInSendToTopic() throws NotificationException {
          AtlasAsyncImportRequest mockImportRequest      = mock(AtlasAsyncImportRequest.class);
          AtlasImportResult       mockResult             = mock(AtlasImportResult.class);
          EntityImportStream      mockEntityImportStream = mock(EntityImportStream.class);
          AtlasEntityWithExtInfo  mockEntityWithExtInfo  = mock(AtlasEntityWithExtInfo.class);
          AtlasEntity             mockEntity             = mock(AtlasEntity.class);
  
          when(mockImportRequest.getTopicName()).thenReturn("test-topic");
          when(mockImportRequest.getImportId()).thenReturn("import-id");
          when(mockImportRequest.getImportDetails()).thenReturn(new AtlasAsyncImportRequest.ImportDetails());
          when(mockImportRequest.getImportTrackingInfo()).thenReturn(new AtlasAsyncImportRequest.ImportTrackingInfo());
          when(mockImportRequest.getImportResult()).thenReturn(mockResult);
          when(mockResult.getUserName()).thenReturn("test-user-1");
          when(mockEntityImportStream.getPosition()).thenReturn(1);
          when(mockEntityImportStream.hasNext()).thenReturn(true, false);
          when(mockEntityImportStream.getNextEntityWithExtInfo()).thenReturn(mockEntityWithExtInfo);
          when(mockEntityWithExtInfo.getEntity()).thenReturn(mockEntity);
          when(mockEntity.getGuid()).thenReturn("entity-guid");
  
          doThrow(new NotificationException(new Exception("Error in sendToTopic")))
                  .when(notificationInterface)
                  .send(eq("test-topic"), anyList(), any());
  
          asyncImportTaskExecutor.publishEntityNotification(mockImportRequest, mockEntityImportStream);
  
          verify(notificationInterface).send(eq("test-topic"), anyList(), any());
          verify(mockEntityImportStream, never()).onImportComplete("entity-guid");
          verify(importService).updateImportRequest(mockImportRequest);
          assertEquals(mockImportRequest.getImportTrackingInfo().getStartEntityPosition(), 1);
          assertEquals(mockImportRequest.getImportDetails().getFailedEntitiesCount(), 1);
          assertEquals(mockImportRequest.getImportDetails().getPublishedEntityCount(), 0);
      }
  
      @Test
      void testPublishEntityNotificationIgnoreFailedEntityAndProcessNext() throws NotificationException {
          AtlasAsyncImportRequest mockImportRequest      = mock(AtlasAsyncImportRequest.class);
          AtlasImportResult       mockResult             = mock(AtlasImportResult.class);
          EntityImportStream      mockEntityImportStream = mock(EntityImportStream.class);
          AtlasEntityWithExtInfo  mockEntityWithExtInfo  = mock(AtlasEntityWithExtInfo.class);
          AtlasEntity             mockEntity             = mock(AtlasEntity.class);
  
          when(mockImportRequest.getTopicName()).thenReturn("test-topic");
          when(mockImportRequest.getImportId()).thenReturn("import-id");
          when(mockImportRequest.getImportDetails()).thenReturn(new AtlasAsyncImportRequest.ImportDetails());
          when(mockImportRequest.getImportTrackingInfo()).thenReturn(new AtlasAsyncImportRequest.ImportTrackingInfo());
          when(mockImportRequest.getImportResult()).thenReturn(mockResult);
          when(mockResult.getUserName()).thenReturn("test-user-1");
          when(mockEntityImportStream.getPosition()).thenReturn(1, 2);
          when(mockEntityImportStream.hasNext()).thenReturn(true, true, false); // Two entities
          when(mockEntityImportStream.getNextEntityWithExtInfo()).thenReturn(mockEntityWithExtInfo);
          when(mockEntityWithExtInfo.getEntity()).thenReturn(mockEntity);
          when(mockEntity.getGuid()).thenReturn("entity-guid");
  
          doThrow(new NotificationException(new Exception("Error in sendToTopic")))
                  .doNothing()
                  .when(notificationInterface)
                  .send(eq("test-topic"), anyList(), any());
  
          asyncImportTaskExecutor.publishEntityNotification(mockImportRequest, mockEntityImportStream);
  
          verify(notificationInterface, times(2)).send(eq("test-topic"), anyList(), any());
          verify(mockEntityImportStream, times(1)).onImportComplete("entity-guid");
          verify(importService, times(2)).updateImportRequest(mockImportRequest);
          assertEquals(mockImportRequest.getImportDetails().getPublishedEntityCount(), 1);
          assertEquals(mockImportRequest.getImportDetails().getFailedEntitiesCount(), 1);
      }
  
      @Test
      void testSkipToPositionHappyPath() {
          AtlasAsyncImportRequest mockImportRequest      = mock(AtlasAsyncImportRequest.class);
          EntityImportStream      mockEntityImportStream = mock(EntityImportStream.class);
  
          when(mockImportRequest.getImportTrackingInfo()).thenReturn(new AtlasAsyncImportRequest.ImportTrackingInfo("", 3));  // Skip to position 3
          when(mockEntityImportStream.hasNext()).thenReturn(true, true, true, true, false); // 4 entities in total
          when(mockEntityImportStream.getPosition()).thenReturn(0, 1, 2, 3);
  
          asyncImportTaskExecutor.skipToStartEntityPosition(mockImportRequest, mockEntityImportStream);
  
          verify(mockEntityImportStream, times(3)).next(); // Skip 3 entities
      }
  
      @Test
      void testSkipToPositionSkipToGreaterThanTotalEntities() {
          AtlasAsyncImportRequest mockImportRequest      = mock(AtlasAsyncImportRequest.class);
          EntityImportStream      mockEntityImportStream = mock(EntityImportStream.class);
  
          when(mockImportRequest.getImportTrackingInfo()).thenReturn(new AtlasAsyncImportRequest.ImportTrackingInfo("", 10)); // Skip to position 10
          when(mockEntityImportStream.hasNext()).thenReturn(true, true, true, false); // 3 entities in total
          when(mockEntityImportStream.getPosition()).thenReturn(0, 1, 2);
  
          asyncImportTaskExecutor.skipToStartEntityPosition(mockImportRequest, mockEntityImportStream);
  
          verify(mockEntityImportStream, times(3)).next(); // Skipped all 3 entities
      }
  
      @Test
      void testSkipToPositionNoEntitiesInStream() {
          AtlasAsyncImportRequest mockImportRequest      = mock(AtlasAsyncImportRequest.class);
          EntityImportStream      mockEntityImportStream = mock(EntityImportStream.class);
  
          when(mockImportRequest.getImportTrackingInfo()).thenReturn(new AtlasAsyncImportRequest.ImportTrackingInfo("", 3)); // Skip to position 3
          when(mockEntityImportStream.hasNext()).thenReturn(false); // No entities in the stream
  
          asyncImportTaskExecutor.skipToStartEntityPosition(mockImportRequest, mockEntityImportStream);
  
          verify(mockEntityImportStream, never()).next(); // No entities to skip
      }
  
      @Test
      void testSkipToPositionSkipToEqualsCurrentPosition() {
          AtlasAsyncImportRequest mockImportRequest      = mock(AtlasAsyncImportRequest.class);
          EntityImportStream      mockEntityImportStream = mock(EntityImportStream.class);
  
          when(mockImportRequest.getImportTrackingInfo()).thenReturn(new AtlasAsyncImportRequest.ImportTrackingInfo("", 2)); // Skip to position 2
          when(mockEntityImportStream.hasNext()).thenReturn(true); // At least one entity in the stream
          when(mockEntityImportStream.getPosition()).thenReturn(2); // Already at position 2
  
          asyncImportTaskExecutor.skipToStartEntityPosition(mockImportRequest, mockEntityImportStream);
  
          verify(mockEntityImportStream, never()).next(); // No entities skipped since current position matches skipTo
      }
  
      @DataProvider(name = "registerRequestScenarios")
      public Object[][] registerRequestScenarios() {
          return new Object[][] {{"null", "NEW"},
                  {"SUCCESSFUL", "NEW"},
                  {"PARTIAL_SUCCESS", "NEW"},
                  {"FAILED", "NEW"},
                  {"ABORTED", "NEW"},
                  {"STAGING", "RESUMED"},
                  {"WAITING", "EXISTING"},
                  {"PROCESSING", "EXISTING"}
          };
      }
  
      @Test(dataProvider = "registerRequestScenarios")
      public void testRegisterRequest(String existingStatus, String expectedOutcome) throws AtlasBaseException {
          AtlasImportResult       mockResult      = mock(AtlasImportResult.class);
          AtlasAsyncImportRequest existingRequest = null;
  
          if (!"null".equals(existingStatus)) {
              existingRequest = mock(AtlasAsyncImportRequest.class);
  
              when(existingRequest.getStatus()).thenReturn(AtlasAsyncImportRequest.ImportStatus.valueOf(existingStatus));
              when(existingRequest.getImportDetails()).thenReturn(new AtlasAsyncImportRequest.ImportDetails());
          }
  
          when(importService.fetchImportRequestByImportId("import-id")).thenReturn(existingRequest);
  
          AtlasAsyncImportRequest result = asyncImportTaskExecutor.registerRequest(mockResult, "import-id", 10, Collections.emptyList());
  
          assertNotNull(result);
  
          if ("NEW".equals(expectedOutcome)) {
              verify(importService).saveImportRequest(any(AtlasAsyncImportRequest.class));
          } else if ("RESUMED".equals(expectedOutcome)) {
              verify(existingRequest).setReceivedTime(any(long.class));
              verify(importService).updateImportRequest(existingRequest);
          } else if ("EXISTING".equals(expectedOutcome)) {
              verify(importService, never()).saveImportRequest(any(AtlasAsyncImportRequest.class));
              verify(importService, never()).updateImportRequest(any(AtlasAsyncImportRequest.class));
          }
      }
  
      @Test
      public void testRegisterRequestThrowsException() throws AtlasBaseException {
          AtlasImportResult       mockResult        = mock(AtlasImportResult.class);
          AtlasAsyncImportRequest mockImportRequest = mock(AtlasAsyncImportRequest.class);
  
          when(mockImportRequest.getStatus()).thenReturn(AtlasAsyncImportRequest.ImportStatus.SUCCESSFUL);
          when(mockImportRequest.getImportDetails()).thenReturn(new AtlasAsyncImportRequest.ImportDetails());
          when(importService.fetchImportRequestByImportId("import-id")).thenReturn(mockImportRequest);
          doThrow(new AtlasBaseException("Some error while saving")).when(importService).saveImportRequest(any(AtlasAsyncImportRequest.class));
  
          AtlasBaseException exception = expectThrows(AtlasBaseException.class, () -> asyncImportTaskExecutor.registerRequest(mockResult, "import-id", 10, Collections.emptyList()));
  
          assertEquals(exception.getAtlasErrorCode(), AtlasErrorCode.IMPORT_REGISTRATION_FAILED);
      }
  
      @Test
      public void testWithRetrySucceedsAfterLockingConflict() throws Exception {
          AtlasImportResult result = mock(AtlasImportResult.class);
          AtlasAsyncImportRequest newRequest = new AtlasAsyncImportRequest(result);
  
          // First call fails with a PermanentLockingException, second succeeds
          doThrow(new RuntimeException(new PermanentLockingException("lock conflict")))
                  .doNothing()
                  .when(importService).saveImportRequest(any(AtlasAsyncImportRequest.class));
  
          when(importService.fetchImportRequestByImportId("import-id")).thenReturn(null);
  
          AtlasAsyncImportRequest response =
                  asyncImportTaskExecutor.registerRequest(result, "import-id", 5, Collections.emptyList());
  
          assertNotNull(response);
          verify(importService, times(2)).saveImportRequest(any(AtlasAsyncImportRequest.class));
      }
  
      @Test
      public void testWithRetryFailsAfterMaxLockingConflicts() throws Exception {
          AtlasImportResult result = mock(AtlasImportResult.class);
  
          // Always fail with PermanentLockingException
          doThrow(new RuntimeException(new PermanentLockingException("lock conflict")))
                  .when(importService).saveImportRequest(any(AtlasAsyncImportRequest.class));
  
          when(importService.fetchImportRequestByImportId("import-id")).thenReturn(null);
  
          AtlasBaseException ex = expectThrows(
                  AtlasBaseException.class,
                  () -> asyncImportTaskExecutor.registerRequest(result, "import-id", 5, Collections.emptyList()));
  
          assertEquals(ex.getAtlasErrorCode(), AtlasErrorCode.IMPORT_REGISTRATION_FAILED);
          verify(importService, times(3)).saveImportRequest(any(AtlasAsyncImportRequest.class));
      }
  
      @Test
      public void testWithRetryFailsOnNonLockingException() throws Exception {
          AtlasImportResult result = mock(AtlasImportResult.class);
  
          // Fail with a different exception
          doThrow(new RuntimeException("Unexpected error"))
                  .when(importService).saveImportRequest(any(AtlasAsyncImportRequest.class));
  
          when(importService.fetchImportRequestByImportId("import-id")).thenReturn(null);
  
          AtlasBaseException ex = expectThrows(
                  AtlasBaseException.class,
                  () -> asyncImportTaskExecutor.registerRequest(result, "import-id", 5, Collections.emptyList()));
  
          assertEquals(ex.getAtlasErrorCode(), AtlasErrorCode.IMPORT_REGISTRATION_FAILED);
      }
  
      @Test
      public void testAbortAsyncImportRequest() throws AtlasBaseException {
          AtlasAsyncImportRequest mockImportRequest = mock(AtlasAsyncImportRequest.class);
  
          when(mockImportRequest.getTopicName()).thenReturn("ATLAS_IMPORT_12344");
          when(importService.abortImport(any(String.class))).thenReturn(mockImportRequest);
          doNothing().when(notificationInterface).deleteTopic(any(NotificationInterface.NotificationType.class), any(String.class));
  
          asyncImportTaskExecutor.abortAsyncImportRequest("12344");
  
          verify(importService, times(1)).abortImport("12344");
          verify(notificationInterface, times(1)).deleteTopic(NotificationInterface.NotificationType.ASYNC_IMPORT, "ATLAS_IMPORT_12344");
      }
  }
  