supernijia
+366
@@ -0,0 +1,366 @@
|
||||
package com.nanri.aiimage.modules.similarasin.service;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import com.nanri.aiimage.common.service.DistributedJobLockService;
|
||||
import com.nanri.aiimage.config.InstanceMetadata;
|
||||
import com.nanri.aiimage.config.SimilarAsinProperties;
|
||||
import com.nanri.aiimage.config.StorageProperties;
|
||||
import com.nanri.aiimage.modules.coze.service.CozeCredentialPoolService;
|
||||
import com.nanri.aiimage.modules.file.service.LocalFileStorageService;
|
||||
import com.nanri.aiimage.modules.file.service.oss.OssStorageService;
|
||||
import com.nanri.aiimage.modules.similarasin.client.SimilarAsinCozeClient;
|
||||
import com.nanri.aiimage.modules.similarasin.mapper.SimilarAsinFilterConditionMapper;
|
||||
import com.nanri.aiimage.modules.similarasin.model.dto.SimilarAsinParsedPayloadDto;
|
||||
import com.nanri.aiimage.modules.similarasin.model.dto.SimilarAsinSourceFileDto;
|
||||
import com.nanri.aiimage.modules.similarasin.model.dto.SimilarAsinSubmitResultRequest;
|
||||
import com.nanri.aiimage.modules.similarasin.model.vo.SimilarAsinParsedRowVo;
|
||||
import com.nanri.aiimage.modules.similarasin.util.SimilarAsinImageEmbedder;
|
||||
import com.nanri.aiimage.modules.task.mapper.FileResultMapper;
|
||||
import com.nanri.aiimage.modules.task.mapper.FileTaskMapper;
|
||||
import com.nanri.aiimage.modules.task.mapper.TaskChunkMapper;
|
||||
import com.nanri.aiimage.modules.task.mapper.TaskScopeStateMapper;
|
||||
import com.nanri.aiimage.modules.task.model.entity.FileResultEntity;
|
||||
import com.nanri.aiimage.modules.task.model.entity.FileTaskEntity;
|
||||
import com.nanri.aiimage.modules.task.model.entity.TaskChunkEntity;
|
||||
import com.nanri.aiimage.modules.task.model.entity.TaskScopeStateEntity;
|
||||
import com.nanri.aiimage.modules.task.service.TaskDistributedLockService;
|
||||
import com.nanri.aiimage.modules.task.service.TaskFileJobService;
|
||||
import com.nanri.aiimage.modules.task.service.TaskProgressSnapshotService;
|
||||
import com.nanri.aiimage.modules.task.service.TransientPayloadStorageService;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeAll;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.extension.ExtendWith;
|
||||
import org.mockito.ArgumentCaptor;
|
||||
import org.mockito.InjectMocks;
|
||||
import org.mockito.Mock;
|
||||
import org.mockito.Spy;
|
||||
import org.mockito.junit.jupiter.MockitoExtension;
|
||||
import org.springframework.transaction.PlatformTransactionManager;
|
||||
import org.springframework.transaction.TransactionDefinition;
|
||||
import org.springframework.transaction.TransactionStatus;
|
||||
import com.baomidou.mybatisplus.core.MybatisConfiguration;
|
||||
import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
|
||||
import org.apache.ibatis.builder.MapperBuilderAssistant;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
import static org.junit.jupiter.api.Assertions.assertFalse;
|
||||
import static org.junit.jupiter.api.Assertions.assertThrows;
|
||||
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.ArgumentMatchers.anyLong;
|
||||
import static org.mockito.ArgumentMatchers.anyString;
|
||||
import static org.mockito.ArgumentMatchers.eq;
|
||||
import static org.mockito.Mockito.doAnswer;
|
||||
import static org.mockito.Mockito.lenient;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.never;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
@ExtendWith(MockitoExtension.class)
|
||||
class SimilarAsinTaskServiceSubmitTest {
|
||||
|
||||
private static final Long TASK_ID = 21879L;
|
||||
private static final String PARSED_POINTER = "rustfs:task-parsed/similar-asin/21879/payload.json";
|
||||
private static final String CHUNK_POINTER = "rustfs:task-chunk/similar-asin/21879/chunk.json";
|
||||
private static final String STORED_CHUNK_POINTER = "\"" + CHUNK_POINTER + "\"";
|
||||
|
||||
@Mock private LocalFileStorageService localFileStorageService;
|
||||
@Mock private OssStorageService ossStorageService;
|
||||
@Mock private StorageProperties storageProperties;
|
||||
@Mock private FileTaskMapper fileTaskMapper;
|
||||
@Mock private FileResultMapper fileResultMapper;
|
||||
@Mock private TaskScopeStateMapper taskScopeStateMapper;
|
||||
@Mock private TaskChunkMapper taskChunkMapper;
|
||||
@Mock private SimilarAsinFilterConditionMapper filterConditionMapper;
|
||||
@Spy private ObjectMapper objectMapper = new ObjectMapper();
|
||||
@Mock private SimilarAsinCozeClient cozeClient;
|
||||
@Mock private SimilarAsinTaskCacheService taskCacheService;
|
||||
@Mock private SimilarAsinProperties properties;
|
||||
@Mock private TaskFileJobService taskFileJobService;
|
||||
@Mock private TaskDistributedLockService taskDistributedLockService;
|
||||
@Mock private TaskProgressSnapshotService taskProgressSnapshotService;
|
||||
@Mock private TransientPayloadStorageService transientPayloadStorageService;
|
||||
@Mock private PlatformTransactionManager transactionManager;
|
||||
@Mock private DistributedJobLockService distributedJobLockService;
|
||||
@Mock private InstanceMetadata instanceMetadata;
|
||||
@Mock private CozeCredentialPoolService cozeCredentialPoolService;
|
||||
@Mock private SimilarAsinImageEmbedder imageEmbedder;
|
||||
@Mock private SimilarAsinImagePrefetchService imagePrefetchService;
|
||||
@Mock private TransactionStatus transactionStatus;
|
||||
|
||||
@InjectMocks private SimilarAsinTaskService service;
|
||||
|
||||
private final AtomicBoolean transactionActive = new AtomicBoolean();
|
||||
|
||||
@BeforeAll
|
||||
static void initializeMybatisMetadata() {
|
||||
MapperBuilderAssistant assistant = new MapperBuilderAssistant(new MybatisConfiguration(), "");
|
||||
TableInfoHelper.initTableInfo(assistant, FileTaskEntity.class);
|
||||
TableInfoHelper.initTableInfo(assistant, FileResultEntity.class);
|
||||
TableInfoHelper.initTableInfo(assistant, TaskChunkEntity.class);
|
||||
TableInfoHelper.initTableInfo(assistant, TaskScopeStateEntity.class);
|
||||
}
|
||||
|
||||
@BeforeEach
|
||||
void setUpTransactionAndLock() {
|
||||
lenient().when(instanceMetadata.getInstanceId()).thenReturn("instance-a");
|
||||
lenient().when(taskDistributedLockService.acquire(
|
||||
eq(SimilarAsinTaskService.MODULE_TYPE),
|
||||
anyLong(),
|
||||
any(Duration.class),
|
||||
eq(10_000L)))
|
||||
.thenReturn(mock(TaskDistributedLockService.LockHandle.class));
|
||||
lenient().when(transactionManager.getTransaction(any(TransactionDefinition.class)))
|
||||
.thenAnswer(invocation -> {
|
||||
transactionActive.set(true);
|
||||
return transactionStatus;
|
||||
});
|
||||
lenient().doAnswer(invocation -> {
|
||||
transactionActive.set(false);
|
||||
return null;
|
||||
}).when(transactionManager).commit(transactionStatus);
|
||||
lenient().doAnswer(invocation -> {
|
||||
transactionActive.set(false);
|
||||
return null;
|
||||
}).when(transactionManager).rollback(transactionStatus);
|
||||
lenient().when(taskChunkMapper.selectList(any())).thenReturn(List.of());
|
||||
lenient().when(fileResultMapper.selectList(any())).thenReturn(List.of());
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
void shutdownExecutors() {
|
||||
service.shutdownAssembleExecutor();
|
||||
}
|
||||
|
||||
@Test
|
||||
void doneCallbackReadsPayloadOnlyBeforeShortTransaction() throws Exception {
|
||||
FileTaskEntity task = runningTask("instance-a");
|
||||
when(fileTaskMapper.selectById(TASK_ID)).thenReturn(task);
|
||||
configureNewChunkAndScope();
|
||||
when(taskChunkMapper.selectCount(any())).thenReturn(1L);
|
||||
|
||||
List<Boolean> storageCallTransactionStates = new ArrayList<>();
|
||||
when(transientPayloadStorageService.resolvePayload(eq(PARSED_POINTER), anyString()))
|
||||
.thenAnswer(invocation -> {
|
||||
storageCallTransactionStates.add(transactionActive.get());
|
||||
return parsedPayloadJson(List.of(new SimilarAsinParsedRowVo(), new SimilarAsinParsedRowVo()));
|
||||
});
|
||||
when(transientPayloadStorageService.storeChunkPayloadVersioned(
|
||||
eq(SimilarAsinTaskService.MODULE_TYPE), eq(TASK_ID), anyString(), eq(0), anyString()))
|
||||
.thenAnswer(invocation -> {
|
||||
storageCallTransactionStates.add(transactionActive.get());
|
||||
return STORED_CHUNK_POINTER;
|
||||
});
|
||||
when(transientPayloadStorageService.wasLastStoreLocalFallback()).thenReturn(false);
|
||||
doAnswer(invocation -> {
|
||||
assertTrue(transactionActive.get());
|
||||
FileResultEntity result = invocation.getArgument(0);
|
||||
result.setId(501L);
|
||||
return 1;
|
||||
}).when(fileResultMapper).insert(any(FileResultEntity.class));
|
||||
doAnswer(invocation -> {
|
||||
storageCallTransactionStates.add(transactionActive.get());
|
||||
return null;
|
||||
}).when(taskCacheService).deleteTaskCache(TASK_ID);
|
||||
|
||||
service.submitResult(TASK_ID, request(true));
|
||||
|
||||
assertFalse(storageCallTransactionStates.isEmpty());
|
||||
assertTrue(storageCallTransactionStates.stream().noneMatch(Boolean::booleanValue));
|
||||
verify(transientPayloadStorageService).resolvePayload(eq(PARSED_POINTER), anyString());
|
||||
ArgumentCaptor<FileResultEntity> resultCaptor = ArgumentCaptor.forClass(FileResultEntity.class);
|
||||
verify(fileResultMapper).insert(resultCaptor.capture());
|
||||
assertEquals(2, resultCaptor.getValue().getRowCount());
|
||||
assertEquals("germany.xlsx", resultCaptor.getValue().getSourceFilename());
|
||||
}
|
||||
|
||||
@Test
|
||||
void unknownCommitOutcomeDoesNotDeletePossiblyCommittedChunk() throws Exception {
|
||||
FileTaskEntity task = runningTask("instance-a");
|
||||
when(fileTaskMapper.selectById(TASK_ID)).thenReturn(task);
|
||||
configureNewChunkAndScope();
|
||||
when(taskChunkMapper.selectCount(any())).thenReturn(0L);
|
||||
configureChunkStore(false);
|
||||
doAnswer(invocation -> {
|
||||
transactionActive.set(false);
|
||||
throw new IllegalStateException("commit ACK lost");
|
||||
}).when(transactionManager).commit(transactionStatus);
|
||||
|
||||
IllegalStateException thrown = assertThrows(IllegalStateException.class,
|
||||
() -> service.submitResult(TASK_ID, request(false)));
|
||||
|
||||
assertEquals("commit ACK lost", thrown.getMessage());
|
||||
verify(transientPayloadStorageService, never()).extractPointer(anyString());
|
||||
verify(transientPayloadStorageService, never()).deletePayloadIfPresent(anyString());
|
||||
verify(taskCacheService, never()).touchTaskHeartbeat(TASK_ID);
|
||||
}
|
||||
|
||||
@Test
|
||||
void duplicateAlreadyPersistedChunkDoesNotRunPayloadCleanup() throws Exception {
|
||||
FileTaskEntity task = runningTask("instance-a");
|
||||
task.setResultJson("{\"ownerInstanceId\":\"instance-a\"}");
|
||||
when(fileTaskMapper.selectById(TASK_ID)).thenReturn(task);
|
||||
when(taskChunkMapper.selectOne(any())).thenReturn(chunk("\"rustfs:winner\""));
|
||||
configureScopeStorage(null);
|
||||
when(taskChunkMapper.selectCount(any())).thenReturn(1L);
|
||||
|
||||
service.submitResult(TASK_ID, request(false));
|
||||
|
||||
verify(transientPayloadStorageService, never()).extractPointer(anyString());
|
||||
verify(transientPayloadStorageService, never()).deletePayloadIfPresent(anyString());
|
||||
}
|
||||
|
||||
@Test
|
||||
void localFallbackOwnerFailureRollsBackChunkAndKeepsCandidate() throws Exception {
|
||||
FileTaskEntity task = runningTask(null);
|
||||
when(fileTaskMapper.selectById(TASK_ID)).thenReturn(task);
|
||||
configureNewChunkAndScope();
|
||||
when(taskChunkMapper.selectCount(any())).thenReturn(1L);
|
||||
configureChunkStore(true);
|
||||
when(fileTaskMapper.updateById(any(FileTaskEntity.class))).thenReturn(0);
|
||||
|
||||
IllegalStateException thrown = assertThrows(IllegalStateException.class,
|
||||
() -> service.submitResult(TASK_ID, request(false)));
|
||||
|
||||
assertTrue(thrown.getMessage().contains("Failed to bind local fallback task owner"));
|
||||
verify(taskChunkMapper).insert(any(TaskChunkEntity.class));
|
||||
verify(transactionManager).rollback(transactionStatus);
|
||||
verify(transactionManager, never()).commit(transactionStatus);
|
||||
verify(transientPayloadStorageService, never()).deletePayloadIfPresent(anyString());
|
||||
}
|
||||
|
||||
@Test
|
||||
void duplicateLocalCandidateIsNotBoundAndReferencedPayloadIsKept() throws Exception {
|
||||
FileTaskEntity task = runningTask(null);
|
||||
when(fileTaskMapper.selectById(TASK_ID)).thenReturn(task);
|
||||
TaskChunkEntity winner = chunk(STORED_CHUNK_POINTER);
|
||||
when(taskChunkMapper.selectOne(any())).thenReturn(null, winner);
|
||||
configureScopeStorage(null);
|
||||
configureChunkStore(true);
|
||||
when(transientPayloadStorageService.extractPointer(STORED_CHUNK_POINTER)).thenReturn(CHUNK_POINTER);
|
||||
when(taskChunkMapper.selectCount(any())).thenReturn(1L);
|
||||
|
||||
service.submitResult(TASK_ID, request(false));
|
||||
|
||||
verify(taskChunkMapper, never()).insert(any(TaskChunkEntity.class));
|
||||
verify(fileTaskMapper, never()).updateById(any(FileTaskEntity.class));
|
||||
verify(transientPayloadStorageService, never()).deletePayloadIfPresent(STORED_CHUNK_POINTER);
|
||||
}
|
||||
|
||||
@Test
|
||||
void repeatedNonFinalCallbackCannotClearCompletedScope() throws Exception {
|
||||
FileTaskEntity task = runningTask("instance-a");
|
||||
task.setResultJson("{\"ownerInstanceId\":\"instance-a\"}");
|
||||
when(fileTaskMapper.selectById(TASK_ID)).thenReturn(task);
|
||||
TaskChunkEntity existing = chunk("\"rustfs:winner\"");
|
||||
when(taskChunkMapper.selectOne(any())).thenReturn(existing);
|
||||
TaskScopeStateEntity scope = scope(1);
|
||||
when(taskScopeStateMapper.selectOne(any())).thenReturn(scope);
|
||||
when(taskChunkMapper.selectCount(any())).thenReturn(1L);
|
||||
|
||||
service.submitResult(TASK_ID, request(false));
|
||||
|
||||
assertEquals(1, scope.getCompleted());
|
||||
verify(taskScopeStateMapper).updateById(scope);
|
||||
verify(transientPayloadStorageService, never()).storeChunkPayloadVersioned(
|
||||
anyString(), anyLong(), anyString(), any(), anyString());
|
||||
}
|
||||
|
||||
private String parsedPayloadJson(List<SimilarAsinParsedRowVo> rows) throws Exception {
|
||||
SimilarAsinParsedPayloadDto payload = new SimilarAsinParsedPayloadDto();
|
||||
payload.setAllItems(rows);
|
||||
payload.setItems(rows);
|
||||
SimilarAsinSourceFileDto sourceFile = new SimilarAsinSourceFileDto();
|
||||
sourceFile.setFileKey("uploads/germany.xlsx");
|
||||
sourceFile.setOriginalFilename("germany.xlsx");
|
||||
payload.setSourceFiles(List.of(sourceFile));
|
||||
return objectMapper.writeValueAsString(payload);
|
||||
}
|
||||
|
||||
private void configureNewChunkAndScope() {
|
||||
AtomicReference<TaskChunkEntity> chunkRef = new AtomicReference<>();
|
||||
when(taskChunkMapper.selectOne(any())).thenAnswer(invocation -> chunkRef.get());
|
||||
doAnswer(invocation -> {
|
||||
TaskChunkEntity chunk = invocation.getArgument(0);
|
||||
chunk.setId(301L);
|
||||
chunkRef.set(chunk);
|
||||
return 1;
|
||||
}).when(taskChunkMapper).insert(any(TaskChunkEntity.class));
|
||||
configureScopeStorage(null);
|
||||
}
|
||||
|
||||
private void configureScopeStorage(TaskScopeStateEntity initial) {
|
||||
AtomicReference<TaskScopeStateEntity> scopeRef = new AtomicReference<>(initial);
|
||||
when(taskScopeStateMapper.selectOne(any())).thenAnswer(invocation -> scopeRef.get());
|
||||
doAnswer(invocation -> {
|
||||
TaskScopeStateEntity scope = invocation.getArgument(0);
|
||||
scope.setId(401L);
|
||||
scopeRef.set(scope);
|
||||
return 1;
|
||||
}).when(taskScopeStateMapper).insert(any(TaskScopeStateEntity.class));
|
||||
}
|
||||
|
||||
private void configureChunkStore(boolean localFallback) {
|
||||
when(transientPayloadStorageService.storeChunkPayloadVersioned(
|
||||
eq(SimilarAsinTaskService.MODULE_TYPE), eq(TASK_ID), anyString(), eq(0), anyString()))
|
||||
.thenReturn(STORED_CHUNK_POINTER);
|
||||
when(transientPayloadStorageService.wasLastStoreLocalFallback()).thenReturn(localFallback);
|
||||
}
|
||||
|
||||
private FileTaskEntity runningTask(String owner) throws Exception {
|
||||
FileTaskEntity task = new FileTaskEntity();
|
||||
task.setId(TASK_ID);
|
||||
task.setModuleType(SimilarAsinTaskService.MODULE_TYPE);
|
||||
task.setStatus("RUNNING");
|
||||
task.setUserId(7L);
|
||||
String resultJson = "{\"parsedPayloadRef\":\"" + PARSED_POINTER + "\"";
|
||||
if (owner != null) {
|
||||
resultJson += ",\"ownerInstanceId\":\"" + owner + "\"";
|
||||
}
|
||||
task.setResultJson(resultJson + "}");
|
||||
return task;
|
||||
}
|
||||
|
||||
private SimilarAsinSubmitResultRequest request(boolean done) {
|
||||
SimilarAsinSubmitResultRequest request = new SimilarAsinSubmitResultRequest();
|
||||
request.setSubmissionId("similar-asin-21879");
|
||||
request.setChunkIndex(0);
|
||||
request.setChunkTotal(1);
|
||||
request.setDone(done);
|
||||
return request;
|
||||
}
|
||||
|
||||
private TaskChunkEntity chunk(String payload) {
|
||||
TaskChunkEntity chunk = new TaskChunkEntity();
|
||||
chunk.setId(301L);
|
||||
chunk.setTaskId(TASK_ID);
|
||||
chunk.setModuleType(SimilarAsinTaskService.MODULE_TYPE);
|
||||
chunk.setScopeHash("existing-scope");
|
||||
chunk.setChunkIndex(0);
|
||||
chunk.setChunkTotal(1);
|
||||
chunk.setPayloadJson(payload);
|
||||
return chunk;
|
||||
}
|
||||
|
||||
private TaskScopeStateEntity scope(int completed) {
|
||||
TaskScopeStateEntity scope = new TaskScopeStateEntity();
|
||||
scope.setId(401L);
|
||||
scope.setTaskId(TASK_ID);
|
||||
scope.setModuleType(SimilarAsinTaskService.MODULE_TYPE);
|
||||
scope.setScopeKey("similar-asin-21879");
|
||||
scope.setScopeHash("existing-scope");
|
||||
scope.setChunkTotal(1);
|
||||
scope.setCompleted(completed);
|
||||
scope.setStateJson("{\"phase\":\"RECEIVED\",\"coze\":\"PENDING\"}");
|
||||
return scope;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user