feat: 完善知识库批量导入与公共 API
- 新增批量异步导入、状态查询、失败重试与中断恢复链路 - 拆分知识库读取、导入、维护权限并完善 Public API 契约 - 补充数据库迁移、管理端交互、接口说明与相关测试
This commit is contained in:
@@ -108,6 +108,26 @@ public class DocumentParseBridgeServiceImplTest {
|
||||
Assert.assertEquals("# demo", taskInfo.getResult().getPreferredText());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证带 PDF 源信息的任务查询会直达 PDF 服务，不试探无关的 PPTX 服务。
|
||||
*/
|
||||
@Test
|
||||
public void shouldQueryPdfTaskAgainstResolvedService() {
|
||||
FakePdfDocumentParseService pdfService = new FakePdfDocumentParseService();
|
||||
pdfService.taskStatusValue = "completed";
|
||||
FakePptxDocumentParseService pptxService = new FakePptxDocumentParseService();
|
||||
DocumentParseBridgeServiceImpl bridgeService =
|
||||
buildBridgeService(pdfService, pptxService, null, pdfService);
|
||||
|
||||
DocumentParseTaskInfo taskInfo = bridgeService.queryTaskInfo("task-1", buildSource());
|
||||
DocumentParsedResult result = bridgeService.queryResult("task-1", buildSource());
|
||||
|
||||
Assert.assertEquals("completed", taskInfo.getStatus());
|
||||
Assert.assertEquals("# demo", result.getPreferredText());
|
||||
Assert.assertEquals(0, pptxService.queryTaskInfoCallCount);
|
||||
Assert.assertEquals(0, pptxService.queryResultCallCount);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证缺少底层服务时抛出稳定错误码。
|
||||
*/
|
||||
@@ -390,6 +410,8 @@ public class DocumentParseBridgeServiceImplTest {
|
||||
private static class FakePptxDocumentParseService implements PptxDocumentParseService {
|
||||
|
||||
private int parseCallCount;
|
||||
private int queryTaskInfoCallCount;
|
||||
private int queryResultCallCount;
|
||||
|
||||
@Override
|
||||
public ParseResponse parse(ParseRequest request) {
|
||||
@@ -415,6 +437,13 @@ public class DocumentParseBridgeServiceImplTest {
|
||||
|
||||
@Override
|
||||
public ParseResponse queryResult(String taskId) {
|
||||
queryResultCallCount++;
|
||||
throw new UnsupportedOperationException();
|
||||
}
|
||||
|
||||
@Override
|
||||
public ParseTaskInfo queryTaskInfo(String taskId) {
|
||||
queryTaskInfoCallCount++;
|
||||
throw new UnsupportedOperationException();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,913 @@
|
||||
package tech.easyflow.ai.documentimport.task;
|
||||
|
||||
import com.mybatisflex.core.query.QueryWrapper;
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
import org.mockito.Mockito;
|
||||
import org.springframework.transaction.support.TransactionSynchronization;
|
||||
import org.springframework.transaction.support.TransactionSynchronizationManager;
|
||||
import org.springframework.web.multipart.MultipartFile;
|
||||
import tech.easyflow.ai.documentimport.DocumentImportBatchDtos;
|
||||
import tech.easyflow.ai.documentimport.DocumentImportBatchRetryResult;
|
||||
import tech.easyflow.ai.documentimport.ImportCallerContext;
|
||||
import tech.easyflow.ai.documentimport.ImportCallerType;
|
||||
import tech.easyflow.ai.entity.DocumentImportBatch;
|
||||
import tech.easyflow.ai.entity.DocumentImportBatchItem;
|
||||
import tech.easyflow.ai.enums.DocumentImportBatchItemStage;
|
||||
import tech.easyflow.ai.enums.DocumentImportBatchItemStatus;
|
||||
import tech.easyflow.ai.enums.DocumentImportBatchStatus;
|
||||
import tech.easyflow.ai.mapper.DocumentImportBatchItemMapper;
|
||||
import tech.easyflow.ai.mapper.DocumentImportBatchMapper;
|
||||
import tech.easyflow.ai.mapper.DocumentMapper;
|
||||
import tech.easyflow.ai.service.DocumentImportBatchItemService;
|
||||
import tech.easyflow.ai.service.DocumentImportBatchService;
|
||||
import tech.easyflow.common.cache.RedisLockExecutor;
|
||||
import tech.easyflow.common.filestorage.FileStorageService;
|
||||
import tech.easyflow.common.filestorage.FileStorageWriteHandle;
|
||||
import tech.easyflow.common.filestorage.FileStorageWriteResult;
|
||||
import tech.easyflow.common.web.exceptions.BusinessException;
|
||||
|
||||
import java.lang.reflect.Field;
|
||||
import java.math.BigInteger;
|
||||
import java.util.Date;
|
||||
import java.util.List;
|
||||
import java.util.Set;
|
||||
|
||||
/**
|
||||
* {@link DocumentImportBatchAppService} 批次启动与重复策略回归测试。
|
||||
*
|
||||
* @author Codex
|
||||
* @since 2026-07-31
|
||||
*/
|
||||
public class DocumentImportBatchAppServiceTest {
|
||||
|
||||
/**
|
||||
* 验证上传领取通过单条多表更新同步刷新批次进度时间。
|
||||
*/
|
||||
@Test
|
||||
public void uploadShouldAtomicallyClaimItemAndTouchBatch() {
|
||||
TestContext context = createContext();
|
||||
DocumentImportBatch batch = context.batchService.getOne(
|
||||
QueryWrapper.create()
|
||||
);
|
||||
batch.setStatus(DocumentImportBatchStatus.UPLOADING.name());
|
||||
DocumentImportBatchItem pending = pendingUploadItem(BigInteger.valueOf(41));
|
||||
DocumentImportBatchItem uploaded = pendingUploadItem(BigInteger.valueOf(41));
|
||||
uploaded.setStatus(DocumentImportBatchItemStatus.UPLOADED.name());
|
||||
uploaded.setFilePath("/stored/demo.docx");
|
||||
FileStorageWriteHandle handle = writeHandle(pending);
|
||||
String locator = handle.encodeLocator();
|
||||
uploaded.setStorageLocator(locator);
|
||||
uploaded.setCleanupPending(false);
|
||||
Mockito.when(context.batchTracker.requireItem(pending.getId()))
|
||||
.thenReturn(pending, uploaded);
|
||||
Mockito.when(context.itemMapper.claimUpload(
|
||||
Mockito.eq(batch.getId()),
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(batch.getKnowledgeId()),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(1);
|
||||
Mockito.when(context.itemMapper.registerUploadWriteIntent(
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(1);
|
||||
Mockito.when(context.batchTracker.completeUpload(
|
||||
pending.getId(),
|
||||
"/stored/demo.docx",
|
||||
locator
|
||||
)).thenReturn(true);
|
||||
FileStorageService storage = Mockito.mock(FileStorageService.class);
|
||||
Mockito.when(storage.prepareRecoverableWrite(
|
||||
Mockito.anyString(),
|
||||
Mockito.anyString()
|
||||
)).thenReturn(handle);
|
||||
Mockito.when(storage.saveRecoverable(
|
||||
Mockito.any(MultipartFile.class),
|
||||
Mockito.eq(handle)
|
||||
)).thenReturn(new FileStorageWriteResult("/stored/demo.docx", locator));
|
||||
setStorageService(context.service, storage);
|
||||
MultipartFile file = uploadFile("demo.docx", pending.getFileSize());
|
||||
|
||||
DocumentImportBatchDtos.ItemResponse response = context.service.uploadItem(
|
||||
batch.getKnowledgeId(),
|
||||
batch.getId(),
|
||||
pending.getId(),
|
||||
file,
|
||||
publicCaller()
|
||||
);
|
||||
|
||||
Assert.assertEquals(
|
||||
DocumentImportBatchItemStatus.UPLOADED.name(),
|
||||
response.getStatus()
|
||||
);
|
||||
org.mockito.InOrder order = Mockito.inOrder(context.itemMapper, storage);
|
||||
order.verify(context.itemMapper).claimUpload(
|
||||
Mockito.eq(batch.getId()),
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(batch.getKnowledgeId()),
|
||||
Mockito.any(Date.class)
|
||||
);
|
||||
order.verify(context.itemMapper).registerUploadWriteIntent(
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.any(Date.class)
|
||||
);
|
||||
order.verify(storage).saveRecoverable(file, handle);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证对象绑定失败后先清理，再恢复为可重新上传状态。
|
||||
*/
|
||||
@Test
|
||||
public void uploadBindFailureShouldCleanupAndAllowSameItemRetry() {
|
||||
TestContext context = createContext();
|
||||
DocumentImportBatch batch = context.batchService.getOne(
|
||||
QueryWrapper.create()
|
||||
);
|
||||
batch.setStatus(DocumentImportBatchStatus.UPLOADING.name());
|
||||
DocumentImportBatchItem pending = pendingUploadItem(BigInteger.valueOf(42));
|
||||
FileStorageWriteHandle handle = writeHandle(pending);
|
||||
String locator = handle.encodeLocator();
|
||||
DocumentImportBatchItem cleanup = pendingUploadItem(pending.getId());
|
||||
cleanup.setStatus(DocumentImportBatchItemStatus.UPLOADING.name());
|
||||
cleanup.setStorageLocator(locator);
|
||||
cleanup.setFilePath("/stored/demo.docx");
|
||||
cleanup.setCleanupPending(true);
|
||||
DocumentImportBatchItem retryPending = pendingUploadItem(pending.getId());
|
||||
DocumentImportBatchItem uploaded = pendingUploadItem(pending.getId());
|
||||
uploaded.setStatus(DocumentImportBatchItemStatus.UPLOADED.name());
|
||||
uploaded.setStorageLocator(locator);
|
||||
uploaded.setFilePath("/stored/demo.docx");
|
||||
Mockito.when(context.batchTracker.requireItem(pending.getId()))
|
||||
.thenReturn(pending, pending, cleanup, retryPending, uploaded);
|
||||
Mockito.when(context.itemMapper.claimUpload(
|
||||
Mockito.eq(batch.getId()),
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(batch.getKnowledgeId()),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(1, 1);
|
||||
Mockito.when(context.itemMapper.registerUploadWriteIntent(
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(1, 1);
|
||||
Mockito.when(context.batchTracker.completeUpload(
|
||||
pending.getId(),
|
||||
"/stored/demo.docx",
|
||||
locator
|
||||
)).thenReturn(false, true);
|
||||
Mockito.when(context.itemMapper.markUploadCleanupPending(
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(1);
|
||||
Mockito.when(context.itemMapper.completeUploadingStorageCleanup(
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.eq("/stored/demo.docx"),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(1);
|
||||
FileStorageService storage = Mockito.mock(FileStorageService.class);
|
||||
Mockito.when(storage.prepareRecoverableWrite(
|
||||
Mockito.anyString(),
|
||||
Mockito.anyString()
|
||||
)).thenReturn(handle);
|
||||
Mockito.when(storage.saveRecoverable(
|
||||
Mockito.any(MultipartFile.class),
|
||||
Mockito.eq(handle)
|
||||
)).thenReturn(new FileStorageWriteResult("/stored/demo.docx", locator));
|
||||
setStorageService(context.service, storage);
|
||||
MultipartFile file = uploadFile("demo.docx", pending.getFileSize());
|
||||
|
||||
try {
|
||||
context.service.uploadItem(
|
||||
batch.getKnowledgeId(),
|
||||
batch.getId(),
|
||||
pending.getId(),
|
||||
file,
|
||||
publicCaller()
|
||||
);
|
||||
Assert.fail("Expected cancelled upload rejection");
|
||||
} catch (BusinessException expected) {
|
||||
Assert.assertTrue(expected.getMessage().contains("上传批次已取消"));
|
||||
}
|
||||
|
||||
DocumentImportBatchDtos.ItemResponse retried = context.service.uploadItem(
|
||||
batch.getKnowledgeId(),
|
||||
batch.getId(),
|
||||
pending.getId(),
|
||||
file,
|
||||
publicCaller()
|
||||
);
|
||||
|
||||
Assert.assertEquals(
|
||||
DocumentImportBatchItemStatus.UPLOADED.name(),
|
||||
retried.getStatus()
|
||||
);
|
||||
org.mockito.InOrder order = Mockito.inOrder(context.itemMapper, storage);
|
||||
order.verify(context.itemMapper).registerUploadWriteIntent(
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.any(Date.class)
|
||||
);
|
||||
order.verify(storage).saveRecoverable(file, handle);
|
||||
order.verify(context.itemMapper).markUploadCleanupPending(
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.any(Date.class)
|
||||
);
|
||||
order.verify(storage).deleteRecoverable(handle);
|
||||
order.verify(context.itemMapper).completeUploadingStorageCleanup(
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.eq("/stored/demo.docx"),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class)
|
||||
);
|
||||
Mockito.verify(context.itemMapper, Mockito.times(2)).claimUpload(
|
||||
Mockito.eq(batch.getId()),
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(batch.getKnowledgeId()),
|
||||
Mockito.any(Date.class)
|
||||
);
|
||||
Mockito.verify(context.batchTracker, Mockito.never()).transitionItem(
|
||||
Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(),
|
||||
Mockito.anyBoolean(), Mockito.anyInt()
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证批次并发取消后，对象清理不会把文件项恢复为待上传。
|
||||
*/
|
||||
@Test
|
||||
public void uploadCleanupShouldPreserveConcurrentCancellation() {
|
||||
TestContext context = createContext();
|
||||
DocumentImportBatch batch = context.batchService.getOne(
|
||||
QueryWrapper.create()
|
||||
);
|
||||
batch.setStatus(DocumentImportBatchStatus.UPLOADING.name());
|
||||
DocumentImportBatchItem pending = pendingUploadItem(BigInteger.valueOf(44));
|
||||
FileStorageWriteHandle handle = writeHandle(pending);
|
||||
String locator = handle.encodeLocator();
|
||||
DocumentImportBatchItem cancelled = pendingUploadItem(pending.getId());
|
||||
cancelled.setStatus(DocumentImportBatchItemStatus.CANCELLED.name());
|
||||
cancelled.setStorageLocator(locator);
|
||||
cancelled.setFilePath("/stored/demo.docx");
|
||||
cancelled.setCleanupPending(true);
|
||||
Mockito.when(context.batchTracker.requireItem(pending.getId()))
|
||||
.thenReturn(pending, pending, cancelled);
|
||||
Mockito.when(context.itemMapper.claimUpload(
|
||||
Mockito.eq(batch.getId()),
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(batch.getKnowledgeId()),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(1);
|
||||
Mockito.when(context.itemMapper.registerUploadWriteIntent(
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(1);
|
||||
Mockito.when(context.batchTracker.completeUpload(
|
||||
pending.getId(),
|
||||
"/stored/demo.docx",
|
||||
locator
|
||||
)).thenReturn(false);
|
||||
Mockito.when(context.itemMapper.markUploadCleanupPending(
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(1);
|
||||
Mockito.when(context.itemMapper.completeCancelledStorageCleanup(
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.eq("/stored/demo.docx"),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(1);
|
||||
FileStorageService storage = Mockito.mock(FileStorageService.class);
|
||||
Mockito.when(storage.prepareRecoverableWrite(
|
||||
Mockito.anyString(),
|
||||
Mockito.anyString()
|
||||
)).thenReturn(handle);
|
||||
Mockito.when(storage.saveRecoverable(
|
||||
Mockito.any(MultipartFile.class),
|
||||
Mockito.eq(handle)
|
||||
)).thenReturn(new FileStorageWriteResult("/stored/demo.docx", locator));
|
||||
setStorageService(context.service, storage);
|
||||
|
||||
try {
|
||||
context.service.uploadItem(
|
||||
batch.getKnowledgeId(),
|
||||
batch.getId(),
|
||||
pending.getId(),
|
||||
uploadFile("demo.docx", pending.getFileSize()),
|
||||
publicCaller()
|
||||
);
|
||||
Assert.fail("Expected cancelled upload rejection");
|
||||
} catch (BusinessException expected) {
|
||||
Assert.assertTrue(expected.getMessage().contains("上传批次已取消"));
|
||||
}
|
||||
|
||||
Mockito.verify(storage).deleteRecoverable(handle);
|
||||
Mockito.verify(context.itemMapper).completeCancelledStorageCleanup(
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.eq("/stored/demo.docx"),
|
||||
Mockito.any(Date.class)
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证对象删除后的数据库恢复失败会保留待办，并由下一轮幂等完成。
|
||||
*/
|
||||
@Test
|
||||
public void cleanupRecoveryFailureShouldRemainRetryable() {
|
||||
TestContext context = createContext();
|
||||
DocumentImportBatchItem cleanup = pendingUploadItem(BigInteger.valueOf(45));
|
||||
FileStorageWriteHandle handle = writeHandle(cleanup);
|
||||
String locator = handle.encodeLocator();
|
||||
cleanup.setStatus(DocumentImportBatchItemStatus.UPLOADING.name());
|
||||
cleanup.setStorageLocator(locator);
|
||||
cleanup.setFilePath("/stored/demo.docx");
|
||||
cleanup.setCleanupPending(true);
|
||||
Mockito.when(context.itemService.list(Mockito.any(QueryWrapper.class)))
|
||||
.thenReturn(List.of(cleanup));
|
||||
Mockito.when(context.itemMapper.completeUploadingStorageCleanup(
|
||||
Mockito.eq(cleanup.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.eq("/stored/demo.docx"),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class)
|
||||
)).thenThrow(new RuntimeException("database unavailable"))
|
||||
.thenReturn(1);
|
||||
FileStorageService storage = Mockito.mock(FileStorageService.class);
|
||||
setStorageService(context.service, storage);
|
||||
|
||||
Assert.assertEquals(0, context.service.cleanupCancelledStoredObjects(10));
|
||||
Assert.assertEquals(1, context.service.cleanupCancelledStoredObjects(10));
|
||||
|
||||
Mockito.verify(storage, Mockito.times(2)).deleteRecoverable(handle);
|
||||
Mockito.verify(context.itemMapper, Mockito.times(2))
|
||||
.completeUploadingStorageCleanup(
|
||||
Mockito.eq(cleanup.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.eq("/stored/demo.docx"),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class)
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证写意图登记已提交但调用抛异常时，写入前撤销会清理定位符。
|
||||
*/
|
||||
@Test
|
||||
public void uploadIntentUnknownCommitShouldAbortBeforePhysicalWrite() {
|
||||
TestContext context = createContext();
|
||||
DocumentImportBatch batch = context.batchService.getOne(
|
||||
QueryWrapper.create()
|
||||
);
|
||||
batch.setStatus(DocumentImportBatchStatus.UPLOADING.name());
|
||||
DocumentImportBatchItem pending = pendingUploadItem(BigInteger.valueOf(46));
|
||||
FileStorageWriteHandle handle = writeHandle(pending);
|
||||
String locator = handle.encodeLocator();
|
||||
Mockito.when(context.batchTracker.requireItem(pending.getId()))
|
||||
.thenReturn(pending);
|
||||
Mockito.when(context.itemMapper.claimUpload(
|
||||
Mockito.eq(batch.getId()),
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(batch.getKnowledgeId()),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(1);
|
||||
Mockito.when(context.itemMapper.registerUploadWriteIntent(
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.any(Date.class)
|
||||
)).thenThrow(new RuntimeException("commit result unknown"));
|
||||
Mockito.when(context.itemMapper.abortUploadBeforeWrite(
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(1);
|
||||
FileStorageService storage = Mockito.mock(FileStorageService.class);
|
||||
Mockito.when(storage.prepareRecoverableWrite(
|
||||
Mockito.anyString(),
|
||||
Mockito.anyString()
|
||||
)).thenReturn(handle);
|
||||
setStorageService(context.service, storage);
|
||||
|
||||
try {
|
||||
context.service.uploadItem(
|
||||
batch.getKnowledgeId(),
|
||||
batch.getId(),
|
||||
pending.getId(),
|
||||
uploadFile("demo.docx", pending.getFileSize()),
|
||||
publicCaller()
|
||||
);
|
||||
Assert.fail("Expected unknown commit exception");
|
||||
} catch (RuntimeException expected) {
|
||||
Assert.assertEquals("commit result unknown", expected.getMessage());
|
||||
}
|
||||
|
||||
Mockito.verify(context.itemMapper).abortUploadBeforeWrite(
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class)
|
||||
);
|
||||
Mockito.verify(storage, Mockito.never()).saveRecoverable(
|
||||
Mockito.any(MultipartFile.class),
|
||||
Mockito.any(FileStorageWriteHandle.class)
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证提交响应异常但数据库已完成上传时不会误删有效对象。
|
||||
*/
|
||||
@Test
|
||||
public void uploadCommitUnknownShouldKeepConfirmedUploadedObject() {
|
||||
TestContext context = createContext();
|
||||
DocumentImportBatch batch = context.batchService.getOne(
|
||||
QueryWrapper.create()
|
||||
);
|
||||
batch.setStatus(DocumentImportBatchStatus.UPLOADING.name());
|
||||
DocumentImportBatchItem pending = pendingUploadItem(BigInteger.valueOf(43));
|
||||
FileStorageWriteHandle handle = writeHandle(pending);
|
||||
String locator = handle.encodeLocator();
|
||||
DocumentImportBatchItem committed = pendingUploadItem(pending.getId());
|
||||
committed.setStatus(DocumentImportBatchItemStatus.UPLOADED.name());
|
||||
committed.setFilePath("/stored/demo.docx");
|
||||
committed.setStorageLocator(locator);
|
||||
committed.setCleanupPending(false);
|
||||
Mockito.when(context.batchTracker.requireItem(pending.getId()))
|
||||
.thenReturn(pending, committed);
|
||||
Mockito.when(context.itemMapper.claimUpload(
|
||||
Mockito.eq(batch.getId()),
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(batch.getKnowledgeId()),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(1);
|
||||
Mockito.when(context.itemMapper.registerUploadWriteIntent(
|
||||
Mockito.eq(pending.getId()),
|
||||
Mockito.eq(locator),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(1);
|
||||
Mockito.when(context.batchTracker.completeUpload(
|
||||
pending.getId(),
|
||||
"/stored/demo.docx",
|
||||
locator
|
||||
)).thenThrow(new RuntimeException("commit result unknown"));
|
||||
FileStorageService storage = Mockito.mock(FileStorageService.class);
|
||||
Mockito.when(storage.prepareRecoverableWrite(
|
||||
Mockito.anyString(),
|
||||
Mockito.anyString()
|
||||
)).thenReturn(handle);
|
||||
Mockito.when(storage.saveRecoverable(
|
||||
Mockito.any(MultipartFile.class),
|
||||
Mockito.eq(handle)
|
||||
)).thenReturn(new FileStorageWriteResult("/stored/demo.docx", locator));
|
||||
setStorageService(context.service, storage);
|
||||
|
||||
DocumentImportBatchDtos.ItemResponse response = context.service.uploadItem(
|
||||
batch.getKnowledgeId(),
|
||||
batch.getId(),
|
||||
pending.getId(),
|
||||
uploadFile("demo.docx", pending.getFileSize()),
|
||||
publicCaller()
|
||||
);
|
||||
|
||||
Assert.assertEquals(
|
||||
DocumentImportBatchItemStatus.UPLOADED.name(),
|
||||
response.getStatus()
|
||||
);
|
||||
Mockito.verify(context.itemMapper, Mockito.never())
|
||||
.markUploadCleanupPending(
|
||||
Mockito.any(),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class)
|
||||
);
|
||||
Mockito.verify(storage, Mockito.never())
|
||||
.deleteRecoverable(Mockito.any(FileStorageWriteHandle.class));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证同一知识库已有运行中的自动批次时拒绝再次启动。
|
||||
*/
|
||||
@Test
|
||||
public void startAutoShouldRejectOtherRunningBatch() {
|
||||
TestContext context = createContext();
|
||||
Mockito.when(context.batchService.count(Mockito.any(QueryWrapper.class))).thenReturn(1L);
|
||||
beginTransactionSynchronization();
|
||||
try {
|
||||
context.service.startBatch(startRequest("AUTO", "SKIP"));
|
||||
Assert.fail("Expected active automatic batch rejection");
|
||||
} catch (BusinessException expected) {
|
||||
Assert.assertTrue(expected.getMessage().contains("已有自动导入批次"));
|
||||
} finally {
|
||||
completeTransactionSynchronization(TransactionSynchronization.STATUS_ROLLED_BACK);
|
||||
}
|
||||
Mockito.verify(context.batchMapper, Mockito.never())
|
||||
.updateByQuery(Mockito.any(DocumentImportBatch.class), Mockito.any());
|
||||
Mockito.verify(context.lockHandle).release();
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证覆盖策略会持久化历史文档 ID，并继续创建新文档任务。
|
||||
*/
|
||||
@Test
|
||||
public void startOverwriteShouldRecordHistoricalDocument() {
|
||||
TestContext context = createContext();
|
||||
Mockito.when(context.batchService.count(Mockito.any(QueryWrapper.class))).thenReturn(0L);
|
||||
DocumentImportBatchItem uploaded = uploadedItem(BigInteger.valueOf(11), BigInteger.ONE);
|
||||
DocumentImportBatchItem historical = uploadedItem(BigInteger.valueOf(21), BigInteger.valueOf(22));
|
||||
historical.setBatchId(BigInteger.valueOf(99));
|
||||
historical.setDocumentId(BigInteger.valueOf(23));
|
||||
historical.setStatus(DocumentImportBatchItemStatus.COMPLETED.name());
|
||||
historical.setCreated(new Date());
|
||||
Mockito.when(context.itemService.list(Mockito.any(QueryWrapper.class)))
|
||||
.thenReturn(List.of(uploaded), List.of(historical));
|
||||
tech.easyflow.ai.entity.Document historicalDocument = new tech.easyflow.ai.entity.Document();
|
||||
historicalDocument.setId(historical.getDocumentId());
|
||||
historicalDocument.setCollectionId(BigInteger.TWO);
|
||||
Mockito.when(context.documentMapper.selectListByQuery(Mockito.any()))
|
||||
.thenReturn(List.of(historicalDocument));
|
||||
Mockito.when(context.batchMapper.updateByQuery(
|
||||
Mockito.any(DocumentImportBatch.class), Mockito.any()
|
||||
)).thenReturn(1);
|
||||
Mockito.when(context.batchTracker.refreshBatch(BigInteger.ONE))
|
||||
.thenReturn(new DocumentImportBatchDtos.StatusResponse());
|
||||
|
||||
beginTransactionSynchronization();
|
||||
try {
|
||||
context.service.startBatch(startRequest("AUTO", "OVERWRITE"));
|
||||
} finally {
|
||||
completeTransactionSynchronization(TransactionSynchronization.STATUS_COMMITTED);
|
||||
}
|
||||
|
||||
Mockito.verify(context.batchTracker)
|
||||
.markReplacement(uploaded.getId(), historical.getDocumentId());
|
||||
Mockito.verify(context.taskAppService)
|
||||
.createBatchImportTasks(Mockito.any(DocumentImportBatch.class), Mockito.anyList());
|
||||
Mockito.verify(context.lockHandle).release();
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证重试代次由服务端读取并原子领取。
|
||||
*/
|
||||
@Test
|
||||
public void retryShouldClaimCurrentGeneration() {
|
||||
TestContext context = createContext();
|
||||
DocumentImportBatch batch = context.batchService.getOne(
|
||||
QueryWrapper.create()
|
||||
);
|
||||
batch.setStatus(DocumentImportBatchStatus.PARTIAL_SUCCEEDED.name());
|
||||
batch.setRetryGeneration(0);
|
||||
DocumentImportBatchItem failed = uploadedItem(
|
||||
BigInteger.valueOf(31),
|
||||
batch.getId()
|
||||
);
|
||||
failed.setStatus(DocumentImportBatchItemStatus.FAILED.name());
|
||||
failed.setRetryable(true);
|
||||
Mockito.when(context.itemService.list(Mockito.any(QueryWrapper.class)))
|
||||
.thenReturn(List.of(failed));
|
||||
Mockito.when(context.batchMapper.claimRetry(
|
||||
Mockito.any(),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(),
|
||||
Mockito.anyInt(),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(1);
|
||||
|
||||
beginTransactionSynchronization();
|
||||
DocumentImportBatchRetryResult first;
|
||||
try {
|
||||
first = context.service.retryOwnedBatch(
|
||||
batch.getId(),
|
||||
publicCaller(),
|
||||
Set.of("file-key")
|
||||
);
|
||||
} finally {
|
||||
completeTransactionSynchronization(
|
||||
TransactionSynchronization.STATUS_ROLLED_BACK
|
||||
);
|
||||
}
|
||||
|
||||
Assert.assertEquals(Integer.valueOf(1), first.getRetryGeneration());
|
||||
Assert.assertEquals(Integer.valueOf(1), first.getRetriedCount());
|
||||
Mockito.verify(context.batchMapper).claimRetry(
|
||||
Mockito.any(),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(),
|
||||
Mockito.eq(0),
|
||||
Mockito.any(Date.class)
|
||||
);
|
||||
Mockito.verify(context.lockHandle).release();
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证修复前的可恢复失败项会在重试领取事务中恢复重试资格和批次计数。
|
||||
*/
|
||||
@Test
|
||||
public void retryShouldRestoreLegacyRecoverableItem() {
|
||||
TestContext context = createContext();
|
||||
DocumentImportBatch batch = context.batchService.getOne(
|
||||
QueryWrapper.create()
|
||||
);
|
||||
batch.setStatus(DocumentImportBatchStatus.PARTIAL_SUCCEEDED.name());
|
||||
batch.setRetryGeneration(0);
|
||||
DocumentImportBatchItem failed = uploadedItem(
|
||||
BigInteger.valueOf(32),
|
||||
batch.getId()
|
||||
);
|
||||
failed.setStage(DocumentImportBatchItemStage.PARSE.name());
|
||||
failed.setStatus(DocumentImportBatchItemStatus.FAILED.name());
|
||||
failed.setRetryable(false);
|
||||
failed.setFailureCode("parse_failed");
|
||||
failed.setErrorSummary("历史代码异常");
|
||||
Mockito.when(context.itemService.list(Mockito.any(QueryWrapper.class)))
|
||||
.thenReturn(List.of(failed));
|
||||
Mockito.when(context.taskAppService.isRecoverableBatchFailure(failed))
|
||||
.thenReturn(true);
|
||||
Mockito.when(context.batchTracker.transitionItem(
|
||||
Mockito.eq(failed.getId()),
|
||||
Mockito.eq(DocumentImportBatchItemStage.PARSE),
|
||||
Mockito.eq(DocumentImportBatchItemStatus.FAILED),
|
||||
Mockito.eq("历史代码异常"),
|
||||
Mockito.eq(true),
|
||||
Mockito.eq(0),
|
||||
Mockito.eq("parse_failed")
|
||||
)).thenReturn(true);
|
||||
Mockito.when(context.batchMapper.claimRetry(
|
||||
Mockito.any(),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(),
|
||||
Mockito.anyInt(),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(1);
|
||||
|
||||
beginTransactionSynchronization();
|
||||
DocumentImportBatchRetryResult result;
|
||||
try {
|
||||
result = context.service.retryOwnedBatch(
|
||||
batch.getId(),
|
||||
publicCaller(),
|
||||
Set.of()
|
||||
);
|
||||
} finally {
|
||||
completeTransactionSynchronization(
|
||||
TransactionSynchronization.STATUS_ROLLED_BACK
|
||||
);
|
||||
}
|
||||
|
||||
Assert.assertEquals(Integer.valueOf(1), result.getRetriedCount());
|
||||
Assert.assertTrue(failed.getRetryable());
|
||||
Mockito.verify(context.batchTracker).transitionItem(
|
||||
failed.getId(),
|
||||
DocumentImportBatchItemStage.PARSE,
|
||||
DocumentImportBatchItemStatus.FAILED,
|
||||
"历史代码异常",
|
||||
true,
|
||||
0,
|
||||
"parse_failed"
|
||||
);
|
||||
Mockito.verify(context.batchMapper).claimRetry(
|
||||
Mockito.eq(batch.getId()),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(),
|
||||
Mockito.eq(0),
|
||||
Mockito.any(Date.class)
|
||||
);
|
||||
Mockito.verify(context.lockHandle).release();
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证运行中的任务拒绝再次领取，调用方可继续查询原 taskId。
|
||||
*/
|
||||
@Test
|
||||
public void retryShouldRejectTaskThatIsAlreadyRunning() {
|
||||
TestContext context = createContext();
|
||||
DocumentImportBatch batch = context.batchService.getOne(
|
||||
QueryWrapper.create()
|
||||
);
|
||||
batch.setStatus(DocumentImportBatchStatus.RUNNING.name());
|
||||
|
||||
beginTransactionSynchronization();
|
||||
try {
|
||||
context.service.retryOwnedBatch(
|
||||
BigInteger.ONE,
|
||||
publicCaller(),
|
||||
Set.of()
|
||||
);
|
||||
Assert.fail("Expected running task rejection");
|
||||
} catch (BusinessException expected) {
|
||||
Assert.assertTrue(
|
||||
expected.getMessage().contains("当前任务状态不允许重试")
|
||||
);
|
||||
} finally {
|
||||
completeTransactionSynchronization(
|
||||
TransactionSynchronization.STATUS_ROLLED_BACK
|
||||
);
|
||||
}
|
||||
Mockito.verify(context.batchMapper, Mockito.never()).claimRetry(
|
||||
Mockito.any(),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(),
|
||||
Mockito.anyInt(),
|
||||
Mockito.any(Date.class)
|
||||
);
|
||||
Mockito.verify(context.lockHandle).release();
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证候选快照过期后若任务已恢复进展，不会继续取消或清理文件项。
|
||||
*/
|
||||
@Test
|
||||
public void staleCancellationShouldStopWhenConditionalClaimMisses() {
|
||||
TestContext context = createContext();
|
||||
Mockito.when(context.batchMapper.claimStaleCancellation(
|
||||
Mockito.any(),
|
||||
Mockito.any(),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(),
|
||||
Mockito.any(Date.class),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(0);
|
||||
|
||||
beginTransactionSynchronization();
|
||||
boolean cancelled;
|
||||
try {
|
||||
cancelled = context.service.cancelStaleBatch(
|
||||
BigInteger.TWO,
|
||||
BigInteger.ONE,
|
||||
publicCaller(),
|
||||
new Date(System.currentTimeMillis() - 30L * 60L * 1000L)
|
||||
);
|
||||
} finally {
|
||||
completeTransactionSynchronization(
|
||||
TransactionSynchronization.STATUS_ROLLED_BACK
|
||||
);
|
||||
}
|
||||
|
||||
Assert.assertFalse(cancelled);
|
||||
Mockito.verify(context.itemService, Mockito.never())
|
||||
.list(Mockito.any(QueryWrapper.class));
|
||||
Mockito.verifyNoInteractions(context.itemMapper);
|
||||
Mockito.verify(context.lockHandle).release();
|
||||
}
|
||||
|
||||
private TestContext createContext() {
|
||||
DocumentImportBatchService batchService = Mockito.mock(DocumentImportBatchService.class);
|
||||
DocumentImportBatchItemService itemService = Mockito.mock(DocumentImportBatchItemService.class);
|
||||
DocumentImportBatchTracker batchTracker = Mockito.mock(DocumentImportBatchTracker.class);
|
||||
KnowledgeDocumentImportTaskAppService taskAppService =
|
||||
Mockito.mock(KnowledgeDocumentImportTaskAppService.class);
|
||||
DocumentImportBatchMapper batchMapper = Mockito.mock(DocumentImportBatchMapper.class);
|
||||
DocumentImportBatchItemMapper itemMapper =
|
||||
Mockito.mock(DocumentImportBatchItemMapper.class);
|
||||
DocumentMapper documentMapper = Mockito.mock(DocumentMapper.class);
|
||||
RedisLockExecutor redisLockExecutor = Mockito.mock(RedisLockExecutor.class);
|
||||
RedisLockExecutor.LockHandle lockHandle = Mockito.mock(RedisLockExecutor.LockHandle.class);
|
||||
Mockito.when(redisLockExecutor.tryAcquire(
|
||||
Mockito.anyString(), Mockito.any(), Mockito.any()
|
||||
)).thenReturn(lockHandle);
|
||||
|
||||
DocumentImportBatch batch = new DocumentImportBatch();
|
||||
batch.setId(BigInteger.ONE);
|
||||
batch.setKnowledgeId(BigInteger.TWO);
|
||||
batch.setStatus(DocumentImportBatchStatus.READY.name());
|
||||
Mockito.when(batchTracker.requireBatch(BigInteger.ONE)).thenReturn(batch);
|
||||
Mockito.when(batchService.getOne(Mockito.any(QueryWrapper.class)))
|
||||
.thenReturn(batch);
|
||||
Mockito.when(batchMapper.selectOwnedForUpdate(
|
||||
Mockito.any(), Mockito.anyString(), Mockito.any()
|
||||
)).thenReturn(batch);
|
||||
|
||||
DocumentImportBatchAppService service = new DocumentImportBatchAppService(
|
||||
batchService,
|
||||
itemService,
|
||||
batchTracker,
|
||||
new DocumentImportBulkProperties(),
|
||||
taskAppService,
|
||||
batchMapper,
|
||||
itemMapper,
|
||||
documentMapper,
|
||||
redisLockExecutor
|
||||
);
|
||||
return new TestContext(
|
||||
service, batchService, itemService, batchTracker,
|
||||
taskAppService, batchMapper, itemMapper, documentMapper, lockHandle
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 创建 Public API 测试调用者。
|
||||
*
|
||||
* @return Public API 调用者
|
||||
*/
|
||||
private ImportCallerContext publicCaller() {
|
||||
return new ImportCallerContext(
|
||||
ImportCallerType.PUBLIC_API,
|
||||
BigInteger.valueOf(77)
|
||||
);
|
||||
}
|
||||
|
||||
private DocumentImportBatchDtos.StartRequest startRequest(String mode, String duplicatePolicy) {
|
||||
DocumentImportBatchDtos.StartRequest request = new DocumentImportBatchDtos.StartRequest();
|
||||
request.setKnowledgeId(BigInteger.TWO);
|
||||
request.setBatchId(BigInteger.ONE);
|
||||
request.setImportMode(mode);
|
||||
request.setDuplicatePolicy(duplicatePolicy);
|
||||
return request;
|
||||
}
|
||||
|
||||
private DocumentImportBatchItem uploadedItem(BigInteger itemId, BigInteger batchId) {
|
||||
DocumentImportBatchItem item = new DocumentImportBatchItem();
|
||||
item.setId(itemId);
|
||||
item.setBatchId(batchId);
|
||||
item.setKnowledgeId(BigInteger.TWO);
|
||||
item.setClientFileKey("file-key");
|
||||
item.setFileName("demo.docx");
|
||||
item.setFilePath("/demo.docx");
|
||||
item.setStatus(DocumentImportBatchItemStatus.UPLOADED.name());
|
||||
return item;
|
||||
}
|
||||
|
||||
private DocumentImportBatchItem pendingUploadItem(BigInteger itemId) {
|
||||
DocumentImportBatchItem item = new DocumentImportBatchItem();
|
||||
item.setId(itemId);
|
||||
item.setBatchId(BigInteger.ONE);
|
||||
item.setKnowledgeId(BigInteger.TWO);
|
||||
item.setClientFileKey("file-key-" + itemId);
|
||||
item.setFileName("demo.docx");
|
||||
item.setRelativePath("demo.docx");
|
||||
item.setFileSize(4L);
|
||||
item.setStage("UPLOAD");
|
||||
item.setStatus(DocumentImportBatchItemStatus.PENDING.name());
|
||||
item.setCleanupPending(false);
|
||||
return item;
|
||||
}
|
||||
|
||||
private FileStorageWriteHandle writeHandle(DocumentImportBatchItem item) {
|
||||
return new FileStorageWriteHandle(
|
||||
"mock-storage",
|
||||
"",
|
||||
"/tmp/easyflow-test",
|
||||
"knowledge-import/" + item.getBatchId() + "/" + item.getId(),
|
||||
item.getId() + ".docx"
|
||||
);
|
||||
}
|
||||
|
||||
private MultipartFile uploadFile(String fileName, long size) {
|
||||
MultipartFile file = Mockito.mock(MultipartFile.class);
|
||||
Mockito.when(file.getOriginalFilename()).thenReturn(fileName);
|
||||
Mockito.when(file.getSize()).thenReturn(size);
|
||||
Mockito.when(file.isEmpty()).thenReturn(false);
|
||||
return file;
|
||||
}
|
||||
|
||||
private void setStorageService(DocumentImportBatchAppService service,
|
||||
FileStorageService storageService) {
|
||||
try {
|
||||
Field field = DocumentImportBatchAppService.class.getDeclaredField(
|
||||
"storageService"
|
||||
);
|
||||
field.setAccessible(true);
|
||||
field.set(service, storageService);
|
||||
} catch (ReflectiveOperationException error) {
|
||||
throw new AssertionError("Failed to inject storage service", error);
|
||||
}
|
||||
}
|
||||
|
||||
private void beginTransactionSynchronization() {
|
||||
TransactionSynchronizationManager.initSynchronization();
|
||||
}
|
||||
|
||||
private void completeTransactionSynchronization(int status) {
|
||||
List<TransactionSynchronization> synchronizations =
|
||||
TransactionSynchronizationManager.getSynchronizations();
|
||||
for (TransactionSynchronization synchronization : synchronizations) {
|
||||
if (status == TransactionSynchronization.STATUS_COMMITTED) {
|
||||
synchronization.afterCommit();
|
||||
}
|
||||
synchronization.afterCompletion(status);
|
||||
}
|
||||
TransactionSynchronizationManager.clearSynchronization();
|
||||
}
|
||||
|
||||
/**
|
||||
* 测试依赖集合。
|
||||
*/
|
||||
private record TestContext(
|
||||
DocumentImportBatchAppService service,
|
||||
DocumentImportBatchService batchService,
|
||||
DocumentImportBatchItemService itemService,
|
||||
DocumentImportBatchTracker batchTracker,
|
||||
KnowledgeDocumentImportTaskAppService taskAppService,
|
||||
DocumentImportBatchMapper batchMapper,
|
||||
DocumentImportBatchItemMapper itemMapper,
|
||||
DocumentMapper documentMapper,
|
||||
RedisLockExecutor.LockHandle lockHandle
|
||||
) {
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,186 @@
|
||||
package tech.easyflow.ai.documentimport.task;
|
||||
|
||||
import org.junit.Test;
|
||||
import tech.easyflow.ai.documentimport.DocumentImportBatchDtos;
|
||||
import tech.easyflow.ai.entity.DocumentImportBatch;
|
||||
import tech.easyflow.ai.entity.DocumentImportBatchItem;
|
||||
import tech.easyflow.ai.enums.DocumentImportBatchItemStage;
|
||||
import tech.easyflow.ai.enums.DocumentImportBatchItemStatus;
|
||||
import tech.easyflow.ai.enums.DocumentImportBatchStatus;
|
||||
import tech.easyflow.ai.enums.DocumentImportMode;
|
||||
import tech.easyflow.ai.mapper.DocumentImportBatchItemMapper;
|
||||
import tech.easyflow.ai.mapper.DocumentImportBatchMapper;
|
||||
import tech.easyflow.ai.service.DocumentImportBatchItemService;
|
||||
import tech.easyflow.ai.service.DocumentImportBatchService;
|
||||
|
||||
import java.math.BigInteger;
|
||||
import java.util.Date;
|
||||
|
||||
import static org.junit.Assert.assertEquals;
|
||||
import static org.junit.Assert.assertFalse;
|
||||
import static org.junit.Assert.assertTrue;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.ArgumentMatchers.anyBoolean;
|
||||
import static org.mockito.ArgumentMatchers.anyInt;
|
||||
import static org.mockito.ArgumentMatchers.eq;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.never;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
/**
|
||||
* 文档批量导入状态汇总测试。
|
||||
*/
|
||||
public class DocumentImportBatchTrackerTest {
|
||||
|
||||
/**
|
||||
* 验证运行批次的完成、处理、失败与等待数量。
|
||||
*/
|
||||
@Test
|
||||
public void shouldAggregateRunningBatchProgress() {
|
||||
DocumentImportBatchService batchService = mock(DocumentImportBatchService.class);
|
||||
DocumentImportBatchItemService itemService = mock(DocumentImportBatchItemService.class);
|
||||
DocumentImportBatchMapper batchMapper = mock(DocumentImportBatchMapper.class);
|
||||
DocumentImportBatchItemMapper itemMapper = mock(DocumentImportBatchItemMapper.class);
|
||||
DocumentImportBatch batch = batch(DocumentImportBatchStatus.RUNNING, 4);
|
||||
batch.setCompletedCount(1);
|
||||
batch.setProcessingCount(1);
|
||||
batch.setFailedCount(1);
|
||||
batch.setPendingCount(1);
|
||||
when(batchService.getById(batch.getId())).thenReturn(batch);
|
||||
|
||||
DocumentImportBatchTracker tracker =
|
||||
new DocumentImportBatchTracker(batchService, itemService, batchMapper, itemMapper);
|
||||
DocumentImportBatchDtos.StatusResponse response = tracker.refreshBatch(batch.getId());
|
||||
|
||||
assertEquals(1, response.getCompletedCount().intValue());
|
||||
assertEquals(1, response.getProcessingCount().intValue());
|
||||
assertEquals(1, response.getFailedCount().intValue());
|
||||
assertEquals(1, response.getPendingCount().intValue());
|
||||
assertEquals(50, response.getProgressPercent().intValue());
|
||||
assertEquals(DocumentImportBatchStatus.RUNNING.name(), response.getStatus());
|
||||
verify(batchService, never()).updateById(batch, false);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证中断批次在汇总失败项后仍保留可继续状态。
|
||||
*/
|
||||
@Test
|
||||
public void shouldPreserveInterruptedStatusUntilUserContinues() {
|
||||
DocumentImportBatchService batchService = mock(DocumentImportBatchService.class);
|
||||
DocumentImportBatchItemService itemService = mock(DocumentImportBatchItemService.class);
|
||||
DocumentImportBatchMapper batchMapper = mock(DocumentImportBatchMapper.class);
|
||||
DocumentImportBatchItemMapper itemMapper = mock(DocumentImportBatchItemMapper.class);
|
||||
DocumentImportBatch batch = batch(DocumentImportBatchStatus.INTERRUPTED, 2);
|
||||
batch.setCompletedCount(1);
|
||||
batch.setFailedCount(1);
|
||||
batch.setRetryableFailedCount(1);
|
||||
when(batchService.getById(batch.getId())).thenReturn(batch);
|
||||
|
||||
DocumentImportBatchTracker tracker =
|
||||
new DocumentImportBatchTracker(batchService, itemService, batchMapper, itemMapper);
|
||||
DocumentImportBatchDtos.StatusResponse response = tracker.refreshBatch(batch.getId());
|
||||
|
||||
assertEquals(DocumentImportBatchStatus.INTERRUPTED.name(), response.getStatus());
|
||||
assertEquals(100, response.getProgressPercent().intValue());
|
||||
assertEquals(1, response.getRetryableFailedCount().intValue());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证失败项重试通过状态 CAS 增量迁移计数。
|
||||
*/
|
||||
@Test
|
||||
public void shouldMoveRetryableFailureBackToPendingAtomically() {
|
||||
DocumentImportBatchService batchService = mock(DocumentImportBatchService.class);
|
||||
DocumentImportBatchItemService itemService = mock(DocumentImportBatchItemService.class);
|
||||
DocumentImportBatchMapper batchMapper = mock(DocumentImportBatchMapper.class);
|
||||
DocumentImportBatchItemMapper itemMapper = mock(DocumentImportBatchItemMapper.class);
|
||||
DocumentImportBatch batch = batch(DocumentImportBatchStatus.RUNNING, 1);
|
||||
batch.setFailedCount(1);
|
||||
batch.setRetryableFailedCount(1);
|
||||
DocumentImportBatchItem item = new DocumentImportBatchItem();
|
||||
item.setId(BigInteger.TEN);
|
||||
item.setBatchId(batch.getId());
|
||||
item.setStage(DocumentImportBatchItemStage.INDEX.name());
|
||||
item.setStatus(DocumentImportBatchItemStatus.FAILED.name());
|
||||
item.setRetryable(true);
|
||||
when(itemService.getById(item.getId())).thenReturn(item);
|
||||
when(itemMapper.transitionStatus(
|
||||
eq(item.getId()),
|
||||
eq(DocumentImportBatchItemStatus.FAILED.name()),
|
||||
eq(DocumentImportBatchItemStage.INDEX.name()),
|
||||
eq(DocumentImportBatchItemStatus.PENDING.name()),
|
||||
eq(null),
|
||||
eq(null),
|
||||
eq(false),
|
||||
eq(1),
|
||||
any(Date.class)
|
||||
)).thenReturn(1);
|
||||
when(batchService.getById(batch.getId())).thenReturn(batch);
|
||||
|
||||
DocumentImportBatchTracker tracker =
|
||||
new DocumentImportBatchTracker(batchService, itemService, batchMapper, itemMapper);
|
||||
|
||||
assertTrue(tracker.updateItem(item.getId(),
|
||||
DocumentImportBatchItemStage.INDEX,
|
||||
DocumentImportBatchItemStatus.PENDING,
|
||||
null));
|
||||
verify(batchMapper).adjustCounters(
|
||||
eq(batch.getId()),
|
||||
eq(0),
|
||||
eq(0),
|
||||
eq(-1),
|
||||
eq(1),
|
||||
eq(0),
|
||||
eq(0),
|
||||
eq(0),
|
||||
eq(-1),
|
||||
any(Date.class)
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证迟到任务不能把已完成文件重新改为处理中。
|
||||
*/
|
||||
@Test
|
||||
public void shouldRejectLateTransitionFromCompletedToRunning() {
|
||||
DocumentImportBatchService batchService = mock(DocumentImportBatchService.class);
|
||||
DocumentImportBatchItemService itemService = mock(DocumentImportBatchItemService.class);
|
||||
DocumentImportBatchMapper batchMapper = mock(DocumentImportBatchMapper.class);
|
||||
DocumentImportBatchItemMapper itemMapper = mock(DocumentImportBatchItemMapper.class);
|
||||
DocumentImportBatchItem item = new DocumentImportBatchItem();
|
||||
item.setId(BigInteger.TEN);
|
||||
item.setBatchId(BigInteger.ONE);
|
||||
item.setStatus(DocumentImportBatchItemStatus.COMPLETED.name());
|
||||
when(itemService.getById(item.getId())).thenReturn(item);
|
||||
|
||||
DocumentImportBatchTracker tracker =
|
||||
new DocumentImportBatchTracker(batchService, itemService, batchMapper, itemMapper);
|
||||
|
||||
assertFalse(tracker.updateItem(item.getId(),
|
||||
DocumentImportBatchItemStage.INDEX,
|
||||
DocumentImportBatchItemStatus.RUNNING,
|
||||
null));
|
||||
verify(itemMapper, never()).transitionStatus(
|
||||
any(), any(), any(), any(), any(), any(), anyBoolean(), anyInt(), any(Date.class));
|
||||
}
|
||||
|
||||
private DocumentImportBatch batch(DocumentImportBatchStatus status, int totalCount) {
|
||||
DocumentImportBatch batch = new DocumentImportBatch();
|
||||
batch.setId(BigInteger.ONE);
|
||||
batch.setKnowledgeId(BigInteger.TWO);
|
||||
batch.setImportMode(DocumentImportMode.AUTO.name());
|
||||
batch.setStatus(status.name());
|
||||
batch.setTotalCount(totalCount);
|
||||
batch.setTotalBytes(100L);
|
||||
batch.setCompletedCount(0);
|
||||
batch.setProcessingCount(0);
|
||||
batch.setFailedCount(0);
|
||||
batch.setPendingCount(0);
|
||||
batch.setUploadedCount(0);
|
||||
batch.setSkippedCount(0);
|
||||
batch.setCancelledCount(0);
|
||||
batch.setRetryableFailedCount(0);
|
||||
return batch;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
package tech.easyflow.ai.documentimport.task;
|
||||
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
import org.mockito.Mockito;
|
||||
import org.springframework.web.multipart.MultipartFile;
|
||||
import tech.easyflow.ai.documentimport.DocumentImportDtos;
|
||||
import tech.easyflow.ai.entity.DocumentChunk;
|
||||
import tech.easyflow.common.filestorage.FileStorageService;
|
||||
|
||||
import java.io.ByteArrayInputStream;
|
||||
import java.lang.reflect.Field;
|
||||
import java.math.BigInteger;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
|
||||
/**
|
||||
* {@link DocumentImportChunkSnapshotService} 持久化恢复测试。
|
||||
*/
|
||||
public class DocumentImportChunkSnapshotServiceTest {
|
||||
|
||||
/**
|
||||
* 验证分块快照可跨缓存写入并从文件存储完整恢复。
|
||||
*
|
||||
* @throws Exception 反射注入或文件读取异常
|
||||
*/
|
||||
@Test
|
||||
public void shouldPersistAndRestoreChunkSnapshot() throws Exception {
|
||||
String storedPath = "http://localhost/snapshots/9-chunks.json";
|
||||
AtomicReference<byte[]> storedBytes = new AtomicReference<byte[]>();
|
||||
FileStorageService storageService = Mockito.mock(FileStorageService.class);
|
||||
Mockito.when(storageService.save(
|
||||
Mockito.any(MultipartFile.class),
|
||||
Mockito.anyString()
|
||||
)).thenAnswer(invocation -> {
|
||||
MultipartFile file = invocation.getArgument(0);
|
||||
storedBytes.set(file.getBytes());
|
||||
return storedPath;
|
||||
});
|
||||
Mockito.when(storageService.readStream(storedPath))
|
||||
.thenAnswer(invocation -> new ByteArrayInputStream(storedBytes.get()));
|
||||
|
||||
DocumentImportChunkSnapshotService service = new DocumentImportChunkSnapshotService();
|
||||
Field storageField = DocumentImportChunkSnapshotService.class
|
||||
.getDeclaredField("storageService");
|
||||
storageField.setAccessible(true);
|
||||
storageField.set(service, storageService);
|
||||
|
||||
DocumentChunk chunk = new DocumentChunk();
|
||||
chunk.setId(BigInteger.valueOf(11));
|
||||
chunk.setDocumentId(BigInteger.valueOf(9));
|
||||
chunk.setDocumentCollectionId(BigInteger.valueOf(7));
|
||||
chunk.setContent("稳定分块");
|
||||
DocumentImportDtos.PreviewSession session = new DocumentImportDtos.PreviewSession();
|
||||
session.setKnowledgeId(BigInteger.valueOf(7));
|
||||
session.setDocumentId(BigInteger.valueOf(9));
|
||||
session.setDocumentChunks(List.of(chunk));
|
||||
|
||||
String path = service.save(session);
|
||||
DocumentImportDtos.PreviewSession restored = service.load(path);
|
||||
|
||||
Assert.assertEquals(storedPath, path);
|
||||
Assert.assertEquals(session.getKnowledgeId(), restored.getKnowledgeId());
|
||||
Assert.assertEquals(session.getDocumentId(), restored.getDocumentId());
|
||||
Assert.assertEquals(1, restored.getDocumentChunks().size());
|
||||
Assert.assertEquals("稳定分块", restored.getDocumentChunks().get(0).getContent());
|
||||
}
|
||||
}
|
||||
@@ -6,11 +6,15 @@ import org.mockito.ArgumentMatchers;
|
||||
import org.mockito.Mockito;
|
||||
import org.springframework.data.redis.core.StringRedisTemplate;
|
||||
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
|
||||
import tech.easyflow.ai.documentimport.DocumentImportKeys;
|
||||
import tech.easyflow.ai.entity.Document;
|
||||
import tech.easyflow.ai.mapper.DocumentMapper;
|
||||
|
||||
import java.lang.reflect.Field;
|
||||
import java.lang.reflect.Method;
|
||||
import java.math.BigInteger;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
|
||||
/**
|
||||
@@ -18,6 +22,33 @@ import java.util.concurrent.atomic.AtomicReference;
|
||||
*/
|
||||
public class DocumentImportTaskStatusStreamServiceTest {
|
||||
|
||||
/**
|
||||
* 验证状态流会携带稳定错误码，供前端进行本地化展示。
|
||||
*
|
||||
* @throws Exception 反射调用异常
|
||||
*/
|
||||
@Test
|
||||
public void buildDocumentPayloadShouldIncludeTaskErrorCode() throws Exception {
|
||||
Document document = new Document();
|
||||
document.setId(BigInteger.valueOf(88));
|
||||
document.setCollectionId(BigInteger.valueOf(99));
|
||||
document.setOptions(new LinkedHashMap<String, Object>(Map.of(
|
||||
DocumentImportKeys.KEY_DOCUMENT_TASK_ERROR_CODE,
|
||||
"parse_service_unavailable"
|
||||
)));
|
||||
DocumentImportTaskStatusStreamService service = new DocumentImportTaskStatusStreamService();
|
||||
Method method = DocumentImportTaskStatusStreamService.class.getDeclaredMethod(
|
||||
"buildDocumentPayload",
|
||||
Document.class
|
||||
);
|
||||
method.setAccessible(true);
|
||||
|
||||
@SuppressWarnings("unchecked")
|
||||
Map<String, Object> payload = (Map<String, Object>) method.invoke(service, document);
|
||||
|
||||
Assert.assertEquals("parse_service_unavailable", payload.get("lastTaskErrorCode"));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证文档状态变更会向 Redis 广播文档 ID。
|
||||
*
|
||||
|
||||
@@ -3,28 +3,53 @@ package tech.easyflow.ai.documentimport.task;
|
||||
import com.easyagents.document.core.entity.DocumentBlock;
|
||||
import com.easyagents.document.core.entity.DocumentImage;
|
||||
import com.easyagents.document.core.entity.DocumentTable;
|
||||
import com.easyagents.core.model.embedding.EmbeddingModel;
|
||||
import com.easyagents.core.store.DocumentStore;
|
||||
import com.easyagents.core.store.StoreOptions;
|
||||
import com.easyagents.core.store.StoreResult;
|
||||
import com.easyagents.rag.ingestion.model.StrategyConfig;
|
||||
import com.easyagents.search.engine.service.DocumentSearcher;
|
||||
import org.apache.ibatis.annotations.Update;
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
import org.mockito.ArgumentCaptor;
|
||||
import org.mockito.Mockito;
|
||||
import org.springframework.web.multipart.MultipartFile;
|
||||
import tech.easyflow.ai.document.exception.DocumentParseBridgeException;
|
||||
import tech.easyflow.ai.document.model.DocumentParseArtifacts;
|
||||
import tech.easyflow.ai.document.model.DocumentParsedResult;
|
||||
import tech.easyflow.ai.document.model.DocumentSourceRef;
|
||||
import tech.easyflow.ai.documentimport.DocumentImportKeys;
|
||||
import tech.easyflow.ai.entity.DocumentChunk;
|
||||
import tech.easyflow.ai.entity.DocumentCollection;
|
||||
import tech.easyflow.ai.entity.DocumentImportBatchItem;
|
||||
import tech.easyflow.ai.entity.DocumentImportTask;
|
||||
import tech.easyflow.ai.enums.DocumentImportBatchItemStage;
|
||||
import tech.easyflow.ai.enums.DocumentImportBatchItemStatus;
|
||||
import tech.easyflow.ai.enums.DocumentImportTaskStatus;
|
||||
import tech.easyflow.ai.enums.DocumentImportTaskPhase;
|
||||
import tech.easyflow.ai.enums.DocumentProcessStatus;
|
||||
import tech.easyflow.ai.mapper.DocumentImportTaskMapper;
|
||||
import tech.easyflow.ai.mapper.DocumentMapper;
|
||||
import tech.easyflow.ai.service.DocumentImportBatchItemService;
|
||||
import tech.easyflow.ai.service.DocumentImportTaskService;
|
||||
import tech.easyflow.ai.service.DocumentService;
|
||||
import tech.easyflow.common.cache.RedisLockExecutor;
|
||||
import tech.easyflow.common.filestorage.FileStorageService;
|
||||
|
||||
import java.io.ByteArrayInputStream;
|
||||
import java.lang.reflect.Constructor;
|
||||
import java.lang.reflect.Field;
|
||||
import java.lang.reflect.Method;
|
||||
import java.lang.reflect.Proxy;
|
||||
import java.math.BigInteger;
|
||||
import java.net.SocketTimeoutException;
|
||||
import java.net.UnknownHostException;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Base64;
|
||||
import java.util.Collection;
|
||||
import java.util.Date;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
@@ -38,6 +63,666 @@ import java.util.concurrent.atomic.AtomicReference;
|
||||
*/
|
||||
public class KnowledgeDocumentImportTaskAppServiceTest {
|
||||
|
||||
/**
|
||||
* 验证待处理任务重新投递只更新必要字段，避免自定义查询结果覆盖非空列。
|
||||
*
|
||||
* @throws Exception 反射注入异常
|
||||
*/
|
||||
@Test
|
||||
public void dispatchPendingTasksShouldTouchPendingWithCas()
|
||||
throws Exception {
|
||||
BigInteger taskId = BigInteger.valueOf(29);
|
||||
DocumentImportTask task = new DocumentImportTask();
|
||||
task.setId(taskId);
|
||||
task.setPhase(DocumentImportTaskPhase.PARSE.name());
|
||||
task.setStatus(DocumentImportTaskStatus.PENDING.name());
|
||||
|
||||
DocumentImportTaskMapper taskMapper =
|
||||
Mockito.mock(DocumentImportTaskMapper.class);
|
||||
Mockito.when(taskMapper.selectPendingFairly(
|
||||
Mockito.any(Date.class), Mockito.anyInt()))
|
||||
.thenReturn(List.of(task));
|
||||
Mockito.when(taskMapper.touchPendingForDispatch(
|
||||
Mockito.eq(taskId), Mockito.any(Date.class), Mockito.any(Date.class),
|
||||
Mockito.any(BigInteger.class)))
|
||||
.thenReturn(1);
|
||||
DocumentImportParseTaskProducer producer =
|
||||
Mockito.mock(DocumentImportParseTaskProducer.class);
|
||||
|
||||
KnowledgeDocumentImportTaskAppService service =
|
||||
new KnowledgeDocumentImportTaskAppService();
|
||||
setField(service, "documentImportTaskMapper", taskMapper);
|
||||
setField(service, "parseTaskProducer", producer);
|
||||
setField(service, "bulkProperties", new DocumentImportBulkProperties());
|
||||
|
||||
service.dispatchPendingTasks();
|
||||
|
||||
ArgumentCaptor<Date> selectCutoffCaptor =
|
||||
ArgumentCaptor.forClass(Date.class);
|
||||
Mockito.verify(taskMapper).selectPendingFairly(
|
||||
selectCutoffCaptor.capture(), Mockito.anyInt());
|
||||
ArgumentCaptor<Date> touchCutoffCaptor =
|
||||
ArgumentCaptor.forClass(Date.class);
|
||||
Mockito.verify(taskMapper).touchPendingForDispatch(
|
||||
Mockito.eq(taskId), touchCutoffCaptor.capture(), Mockito.any(Date.class),
|
||||
Mockito.any(BigInteger.class));
|
||||
Assert.assertEquals(selectCutoffCaptor.getValue(), touchCutoffCaptor.getValue());
|
||||
Mockito.verify(producer).send(taskId);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证任务已离开待处理状态时不再发送重复消息。
|
||||
*
|
||||
* @throws Exception 反射注入异常
|
||||
*/
|
||||
@Test
|
||||
public void dispatchPendingTasksShouldSkipChangedTask()
|
||||
throws Exception {
|
||||
BigInteger taskId = BigInteger.valueOf(30);
|
||||
DocumentImportTask task = new DocumentImportTask();
|
||||
task.setId(taskId);
|
||||
task.setPhase(DocumentImportTaskPhase.PARSE.name());
|
||||
|
||||
DocumentImportTaskMapper taskMapper =
|
||||
Mockito.mock(DocumentImportTaskMapper.class);
|
||||
Mockito.when(taskMapper.selectPendingFairly(
|
||||
Mockito.any(Date.class), Mockito.anyInt()))
|
||||
.thenReturn(List.of(task));
|
||||
Mockito.when(taskMapper.touchPendingForDispatch(
|
||||
Mockito.eq(taskId), Mockito.any(Date.class), Mockito.any(Date.class),
|
||||
Mockito.any(BigInteger.class)))
|
||||
.thenReturn(0);
|
||||
DocumentImportParseTaskProducer producer =
|
||||
Mockito.mock(DocumentImportParseTaskProducer.class);
|
||||
|
||||
KnowledgeDocumentImportTaskAppService service =
|
||||
new KnowledgeDocumentImportTaskAppService();
|
||||
setField(service, "documentImportTaskMapper", taskMapper);
|
||||
setField(service, "parseTaskProducer", producer);
|
||||
setField(service, "bulkProperties", new DocumentImportBulkProperties());
|
||||
|
||||
service.dispatchPendingTasks();
|
||||
|
||||
Mockito.verify(producer, Mockito.never()).send(Mockito.any());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证重新投递资格更新包含待处理状态和过期时间边界，确保竞争更新仅一个成功。
|
||||
*
|
||||
* @throws Exception 映射方法不存在时抛出
|
||||
*/
|
||||
@Test
|
||||
public void touchPendingForDispatchSqlShouldFenceConcurrentDispatch()
|
||||
throws Exception {
|
||||
Method method = DocumentImportTaskMapper.class.getMethod(
|
||||
"touchPendingForDispatch",
|
||||
BigInteger.class,
|
||||
Date.class,
|
||||
Date.class,
|
||||
BigInteger.class
|
||||
);
|
||||
Update update = method.getAnnotation(Update.class);
|
||||
String sql = String.join(" ", update.value());
|
||||
|
||||
Assert.assertTrue(sql.contains("status='PENDING'"));
|
||||
Assert.assertTrue(sql.contains("modified <= #{redispatchBefore}"));
|
||||
Assert.assertTrue(sql.contains("modified=#{now}"));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证分块策略快照使用忽略 null 的部分实体更新。
|
||||
*
|
||||
* @throws Exception 反射调用异常
|
||||
*/
|
||||
@Test
|
||||
public void persistAppliedStrategyShouldKeepRequiredBatchItemFields()
|
||||
throws Exception {
|
||||
BigInteger itemId = BigInteger.valueOf(31);
|
||||
DocumentImportBatchItemService itemService =
|
||||
Mockito.mock(DocumentImportBatchItemService.class);
|
||||
Mockito.when(itemService.updateById(Mockito.any()))
|
||||
.thenReturn(true);
|
||||
KnowledgeDocumentImportTaskAppService service =
|
||||
new KnowledgeDocumentImportTaskAppService();
|
||||
setField(service, "documentImportBatchItemService", itemService);
|
||||
|
||||
Method method =
|
||||
KnowledgeDocumentImportTaskAppService.class.getDeclaredMethod(
|
||||
"persistAppliedStrategy",
|
||||
BigInteger.class,
|
||||
StrategyConfig.class
|
||||
);
|
||||
method.setAccessible(true);
|
||||
method.invoke(service, itemId, StrategyConfig.defaults());
|
||||
|
||||
ArgumentCaptor<DocumentImportBatchItem> updateCaptor =
|
||||
ArgumentCaptor.forClass(DocumentImportBatchItem.class);
|
||||
Mockito.verify(itemService).updateById(updateCaptor.capture());
|
||||
Mockito.verify(itemService, Mockito.never()).updateById(
|
||||
Mockito.any(), Mockito.anyBoolean());
|
||||
Assert.assertEquals(itemId, updateCaptor.getValue().getId());
|
||||
Assert.assertNotNull(
|
||||
updateCaptor.getValue().getStrategySnapshotJson());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证 SPLIT 待处理任务使用 execution token 和租约原子领取。
|
||||
*
|
||||
* @throws Exception 反射注入异常
|
||||
*/
|
||||
@Test
|
||||
public void tryMarkSplitTaskRunningShouldCreateExecutionFence()
|
||||
throws Exception {
|
||||
BigInteger taskId = BigInteger.valueOf(35);
|
||||
DocumentImportTask task = new DocumentImportTask();
|
||||
task.setId(taskId);
|
||||
task.setPhase(DocumentImportTaskPhase.SPLIT.name());
|
||||
task.setStatus(DocumentImportTaskStatus.PENDING.name());
|
||||
|
||||
DocumentImportTaskService taskService =
|
||||
Mockito.mock(DocumentImportTaskService.class);
|
||||
Mockito.when(taskService.getById(taskId)).thenReturn(task);
|
||||
Mockito.when(taskService.count(
|
||||
Mockito.any(com.mybatisflex.core.query.QueryWrapper.class)))
|
||||
.thenReturn(0L);
|
||||
DocumentImportTaskMapper taskMapper =
|
||||
Mockito.mock(DocumentImportTaskMapper.class);
|
||||
Mockito.when(taskMapper.claimPending(
|
||||
Mockito.eq(taskId),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class),
|
||||
Mockito.any(Date.class),
|
||||
Mockito.nullable(BigInteger.class)
|
||||
)).thenReturn(1);
|
||||
RedisLockExecutor lockExecutor =
|
||||
Mockito.mock(RedisLockExecutor.class);
|
||||
RedisLockExecutor.LockHandle lockHandle =
|
||||
Mockito.mock(RedisLockExecutor.LockHandle.class);
|
||||
Mockito.when(lockExecutor.tryAcquire(
|
||||
Mockito.anyString(), Mockito.any(), Mockito.any()
|
||||
)).thenReturn(lockHandle);
|
||||
|
||||
KnowledgeDocumentImportTaskAppService service =
|
||||
new KnowledgeDocumentImportTaskAppService();
|
||||
setField(service, "documentImportTaskService", taskService);
|
||||
setField(service, "documentImportTaskMapper", taskMapper);
|
||||
setField(service, "redisLockExecutor", lockExecutor);
|
||||
setField(service, "bulkProperties", new DocumentImportBulkProperties());
|
||||
|
||||
Assert.assertTrue(service.tryMarkTaskRunning(taskId));
|
||||
|
||||
ArgumentCaptor<String> token = ArgumentCaptor.forClass(String.class);
|
||||
ArgumentCaptor<Date> lease = ArgumentCaptor.forClass(Date.class);
|
||||
Mockito.verify(taskMapper).claimPending(
|
||||
Mockito.eq(taskId),
|
||||
token.capture(),
|
||||
lease.capture(),
|
||||
Mockito.any(Date.class),
|
||||
Mockito.nullable(BigInteger.class)
|
||||
);
|
||||
Assert.assertFalse(token.getValue().isBlank());
|
||||
Assert.assertTrue(lease.getValue().after(new Date()));
|
||||
Mockito.verify(lockHandle).release();
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证新文档完成后会清理待覆盖的历史文档并清除持久化待办。
|
||||
*
|
||||
* @throws Exception 反射注入异常
|
||||
*/
|
||||
@Test
|
||||
public void cleanupCompletedReplacementShouldDeleteHistoricalDocument() throws Exception {
|
||||
BigInteger itemId = BigInteger.valueOf(901);
|
||||
BigInteger documentId = BigInteger.valueOf(902);
|
||||
BigInteger replacedDocumentId = BigInteger.valueOf(903);
|
||||
BigInteger knowledgeId = BigInteger.valueOf(904);
|
||||
|
||||
DocumentImportBatchItem item = new DocumentImportBatchItem();
|
||||
item.setId(itemId);
|
||||
item.setDocumentId(documentId);
|
||||
item.setReplacedDocumentId(replacedDocumentId);
|
||||
tech.easyflow.ai.entity.Document replacement = new tech.easyflow.ai.entity.Document();
|
||||
replacement.setId(documentId);
|
||||
replacement.setCollectionId(knowledgeId);
|
||||
replacement.setProcessStatus(DocumentProcessStatus.COMPLETED.name());
|
||||
tech.easyflow.ai.entity.Document historical = new tech.easyflow.ai.entity.Document();
|
||||
historical.setId(replacedDocumentId);
|
||||
historical.setCollectionId(knowledgeId);
|
||||
historical.setProcessStatus(DocumentProcessStatus.COMPLETED.name());
|
||||
|
||||
DocumentImportBatchTracker tracker = Mockito.mock(DocumentImportBatchTracker.class);
|
||||
Mockito.when(tracker.requireItem(itemId)).thenReturn(item);
|
||||
DocumentMapper documentMapper = Mockito.mock(DocumentMapper.class);
|
||||
Mockito.when(documentMapper.selectOneById(documentId)).thenReturn(replacement);
|
||||
Mockito.when(documentMapper.selectOneById(replacedDocumentId)).thenReturn(historical);
|
||||
DocumentService documentService = Mockito.mock(DocumentService.class);
|
||||
Mockito.when(documentService.removeDoc(replacedDocumentId.toString())).thenReturn(true);
|
||||
RedisLockExecutor redisLockExecutor = Mockito.mock(RedisLockExecutor.class);
|
||||
RedisLockExecutor.LockHandle lockHandle = Mockito.mock(RedisLockExecutor.LockHandle.class);
|
||||
Mockito.when(redisLockExecutor.tryAcquire(
|
||||
Mockito.anyString(), Mockito.any(), Mockito.any()
|
||||
)).thenReturn(lockHandle);
|
||||
|
||||
KnowledgeDocumentImportTaskAppService service = new KnowledgeDocumentImportTaskAppService();
|
||||
setField(service, "documentImportBatchTracker", tracker);
|
||||
setField(service, "documentMapper", documentMapper);
|
||||
setField(service, "documentService", documentService);
|
||||
setField(service, "redisLockExecutor", redisLockExecutor);
|
||||
|
||||
service.cleanupCompletedReplacement(itemId);
|
||||
|
||||
Mockito.verify(documentService).removeDoc(replacedDocumentId.toString());
|
||||
Mockito.verify(documentMapper).deleteById(replacedDocumentId);
|
||||
Mockito.verify(tracker).clearReplacement(itemId, replacedDocumentId);
|
||||
Mockito.verify(lockHandle).release();
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证失联且关联文档已删除的运行任务会被收口，避免持续占用并发名额。
|
||||
*
|
||||
* @throws Exception 反射注入异常
|
||||
*/
|
||||
@Test
|
||||
public void recoverInterruptedTaskShouldFailOrphanWithoutDocument() throws Exception {
|
||||
BigInteger taskId = BigInteger.valueOf(41);
|
||||
Date cutoff = new Date(System.currentTimeMillis() - 60_000L);
|
||||
DocumentImportTask task = new DocumentImportTask();
|
||||
task.setId(taskId);
|
||||
task.setDocumentId(BigInteger.valueOf(42));
|
||||
task.setPhase(DocumentImportTaskPhase.PARSE.name());
|
||||
task.setStatus(DocumentImportTaskStatus.RUNNING.name());
|
||||
task.setModified(new Date(cutoff.getTime() - 1_000L));
|
||||
|
||||
DocumentImportTaskService taskService = Mockito.mock(DocumentImportTaskService.class);
|
||||
Mockito.when(taskService.getById(taskId)).thenReturn(task);
|
||||
DocumentImportTaskMapper taskMapper = Mockito.mock(DocumentImportTaskMapper.class);
|
||||
Mockito.when(taskMapper.failExpiredOwned(
|
||||
Mockito.eq(taskId),
|
||||
Mockito.isNull(),
|
||||
Mockito.anyString(),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class),
|
||||
Mockito.eq(cutoff),
|
||||
Mockito.any()
|
||||
)).thenReturn(1);
|
||||
DocumentMapper documentMapper = Mockito.mock(DocumentMapper.class);
|
||||
Mockito.when(documentMapper.selectOneById(task.getDocumentId())).thenReturn(null);
|
||||
|
||||
KnowledgeDocumentImportTaskAppService service = new KnowledgeDocumentImportTaskAppService();
|
||||
setField(service, "documentImportTaskService", taskService);
|
||||
setField(service, "documentImportTaskMapper", taskMapper);
|
||||
setField(service, "documentMapper", documentMapper);
|
||||
|
||||
service.recoverInterruptedTask(taskId, cutoff);
|
||||
|
||||
Mockito.verify(taskMapper).failExpiredOwned(
|
||||
Mockito.eq(taskId),
|
||||
Mockito.isNull(),
|
||||
Mockito.eq("任务执行中断，请重试"),
|
||||
Mockito.eq("execution_interrupted"),
|
||||
Mockito.any(Date.class),
|
||||
Mockito.eq(cutoff),
|
||||
Mockito.any()
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证运行中断收口后仍保留批次文件的人工重试资格。
|
||||
*
|
||||
* @throws Exception 反射调用异常
|
||||
*/
|
||||
@Test
|
||||
public void recoveredBatchStateShouldRemainRetryableBeforeAttemptLimit()
|
||||
throws Exception {
|
||||
BigInteger batchId = BigInteger.valueOf(61);
|
||||
BigInteger itemId = BigInteger.valueOf(62);
|
||||
DocumentImportTask task = new DocumentImportTask();
|
||||
task.setBatchId(batchId);
|
||||
task.setBatchItemId(itemId);
|
||||
DocumentImportBatchItem item = new DocumentImportBatchItem();
|
||||
item.setId(itemId);
|
||||
item.setAttemptCount(0);
|
||||
|
||||
DocumentImportBatchTracker tracker =
|
||||
Mockito.mock(DocumentImportBatchTracker.class);
|
||||
Mockito.when(tracker.requireItem(itemId)).thenReturn(item);
|
||||
KnowledgeDocumentImportTaskAppService service =
|
||||
new KnowledgeDocumentImportTaskAppService();
|
||||
setField(service, "documentImportBatchTracker", tracker);
|
||||
setField(service, "bulkProperties", new DocumentImportBulkProperties());
|
||||
Method method = KnowledgeDocumentImportTaskAppService.class
|
||||
.getDeclaredMethod(
|
||||
"finishRecoveredBatchState",
|
||||
DocumentImportTask.class,
|
||||
DocumentImportBatchItemStage.class,
|
||||
String.class,
|
||||
String.class
|
||||
);
|
||||
method.setAccessible(true);
|
||||
|
||||
method.invoke(
|
||||
service,
|
||||
task,
|
||||
DocumentImportBatchItemStage.PARSE,
|
||||
"任务执行中断，请继续批次",
|
||||
"execution_interrupted"
|
||||
);
|
||||
|
||||
Mockito.verify(tracker).transitionItem(
|
||||
itemId,
|
||||
DocumentImportBatchItemStage.PARSE,
|
||||
DocumentImportBatchItemStatus.FAILED,
|
||||
"任务执行中断，请继续批次",
|
||||
true,
|
||||
0,
|
||||
"execution_interrupted"
|
||||
);
|
||||
Mockito.verify(tracker).markInterrupted(batchId);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证待处理超时按创建时间判定，即使重新投递刷新了修改时间也会正常收口。
|
||||
*
|
||||
* @throws Exception 反射注入异常
|
||||
*/
|
||||
@Test
|
||||
public void expireTimedOutPendingTaskShouldUseCreatedTime() throws Exception {
|
||||
BigInteger taskId = BigInteger.valueOf(51);
|
||||
BigInteger documentId = BigInteger.valueOf(52);
|
||||
Date cutoff = new Date(System.currentTimeMillis() - 60_000L);
|
||||
DocumentImportTask task = new DocumentImportTask();
|
||||
task.setId(taskId);
|
||||
task.setDocumentId(documentId);
|
||||
task.setPhase(DocumentImportTaskPhase.PARSE.name());
|
||||
task.setStatus(DocumentImportTaskStatus.PENDING.name());
|
||||
task.setCreated(new Date(cutoff.getTime() - 1_000L));
|
||||
task.setModified(new Date());
|
||||
|
||||
tech.easyflow.ai.entity.Document document = new tech.easyflow.ai.entity.Document();
|
||||
document.setId(documentId);
|
||||
document.setCollectionId(BigInteger.valueOf(53));
|
||||
AtomicReference<tech.easyflow.ai.entity.Document> updatedDocumentRef =
|
||||
new AtomicReference<tech.easyflow.ai.entity.Document>();
|
||||
|
||||
DocumentImportTaskService taskService = Mockito.mock(DocumentImportTaskService.class);
|
||||
Mockito.when(taskService.getById(taskId)).thenReturn(task);
|
||||
DocumentImportTaskMapper taskMapper = Mockito.mock(DocumentImportTaskMapper.class);
|
||||
Mockito.when(taskMapper.updateByQuery(
|
||||
Mockito.any(DocumentImportTask.class),
|
||||
Mockito.any()
|
||||
)).thenReturn(1);
|
||||
|
||||
KnowledgeDocumentImportTaskAppService service = new KnowledgeDocumentImportTaskAppService();
|
||||
setField(service, "documentImportTaskService", taskService);
|
||||
setField(service, "documentImportTaskMapper", taskMapper);
|
||||
setField(service, "documentMapper", mockDocumentMapper(document, updatedDocumentRef));
|
||||
setField(service, "documentImportTaskStatusStreamService", new NoopTaskStatusStreamService());
|
||||
|
||||
service.expireTimedOutPendingTask(taskId, cutoff);
|
||||
|
||||
tech.easyflow.ai.entity.Document updatedDocument = updatedDocumentRef.get();
|
||||
Assert.assertNotNull(updatedDocument);
|
||||
Assert.assertEquals(DocumentProcessStatus.PARSE_FAILED.name(), updatedDocument.getProcessStatus());
|
||||
Assert.assertEquals("任务排队超时，请重试", updatedDocument.getLastTaskError());
|
||||
Assert.assertEquals("pending_timeout",
|
||||
updatedDocument.getOptions().get(DocumentImportKeys.KEY_DOCUMENT_TASK_ERROR_CODE));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证未取得 MinerU 任务 ID 的超时提交会失败并释放运行状态。
|
||||
*
|
||||
* @throws Exception 反射注入异常
|
||||
*/
|
||||
@Test
|
||||
public void expireTimedOutParseSubmissionShouldExposeFriendlyTimeout() throws Exception {
|
||||
BigInteger taskId = BigInteger.valueOf(61);
|
||||
BigInteger documentId = BigInteger.valueOf(62);
|
||||
Date cutoff = new Date(System.currentTimeMillis() - 60_000L);
|
||||
DocumentImportTask task = new DocumentImportTask();
|
||||
task.setId(taskId);
|
||||
task.setDocumentId(documentId);
|
||||
task.setPhase(DocumentImportTaskPhase.PARSE.name());
|
||||
task.setStatus(DocumentImportTaskStatus.RUNNING.name());
|
||||
task.setStartedAt(new Date(cutoff.getTime() - 1_000L));
|
||||
|
||||
tech.easyflow.ai.entity.Document document = new tech.easyflow.ai.entity.Document();
|
||||
document.setId(documentId);
|
||||
document.setCollectionId(BigInteger.valueOf(63));
|
||||
AtomicReference<tech.easyflow.ai.entity.Document> updatedDocumentRef =
|
||||
new AtomicReference<tech.easyflow.ai.entity.Document>();
|
||||
|
||||
DocumentImportTaskService taskService = Mockito.mock(DocumentImportTaskService.class);
|
||||
Mockito.when(taskService.getById(taskId)).thenReturn(task);
|
||||
DocumentImportTaskMapper taskMapper = Mockito.mock(DocumentImportTaskMapper.class);
|
||||
Mockito.when(taskMapper.updateByQuery(
|
||||
Mockito.any(DocumentImportTask.class),
|
||||
Mockito.any()
|
||||
)).thenReturn(1);
|
||||
|
||||
KnowledgeDocumentImportTaskAppService service = new KnowledgeDocumentImportTaskAppService();
|
||||
setField(service, "documentImportTaskService", taskService);
|
||||
setField(service, "documentImportTaskMapper", taskMapper);
|
||||
setField(service, "documentMapper", mockDocumentMapper(document, updatedDocumentRef));
|
||||
setField(service, "documentImportTaskStatusStreamService", new NoopTaskStatusStreamService());
|
||||
|
||||
service.expireTimedOutParseSubmission(taskId, cutoff);
|
||||
|
||||
ArgumentCaptor<DocumentImportTask> updateCaptor = ArgumentCaptor.forClass(DocumentImportTask.class);
|
||||
Mockito.verify(taskMapper).updateByQuery(updateCaptor.capture(), Mockito.any());
|
||||
Assert.assertEquals(DocumentImportTaskStatus.FAILED.name(), updateCaptor.getValue().getStatus());
|
||||
Assert.assertEquals("文档解析服务响应超时，请重试", updateCaptor.getValue().getErrorSummary());
|
||||
|
||||
tech.easyflow.ai.entity.Document updatedDocument = updatedDocumentRef.get();
|
||||
Assert.assertNotNull(updatedDocument);
|
||||
Assert.assertEquals(DocumentProcessStatus.PARSE_FAILED.name(), updatedDocument.getProcessStatus());
|
||||
Assert.assertEquals("文档解析服务响应超时，请重试", updatedDocument.getLastTaskError());
|
||||
Assert.assertEquals("parse_service_timeout",
|
||||
updatedDocument.getOptions().get(DocumentImportKeys.KEY_DOCUMENT_TASK_ERROR_CODE));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证 MinerU 任意 5xx 会归一化为服务暂不可用错误。
|
||||
*
|
||||
* @throws Exception 反射调用异常
|
||||
*/
|
||||
@Test
|
||||
public void resolveParseFailureCodeShouldRecognizeMineruServerErrors() throws Exception {
|
||||
KnowledgeDocumentImportTaskAppService service = new KnowledgeDocumentImportTaskAppService();
|
||||
Method method = KnowledgeDocumentImportTaskAppService.class.getDeclaredMethod(
|
||||
"resolveParseFailureCode",
|
||||
Throwable.class
|
||||
);
|
||||
method.setAccessible(true);
|
||||
|
||||
for (int statusCode : new int[] {500, 502, 503, 504, 599}) {
|
||||
Object code = method.invoke(service,
|
||||
new RuntimeException("MinerU request failed: path=/tasks, status=" + statusCode + ", body="));
|
||||
Assert.assertEquals("status=" + statusCode, "parse_service_unavailable", code);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证底层读取超时会归一化为解析服务超时错误。
|
||||
*
|
||||
* @throws Exception 反射调用异常
|
||||
*/
|
||||
@Test
|
||||
public void resolveParseFailureCodeShouldRecognizeSocketTimeout() throws Exception {
|
||||
KnowledgeDocumentImportTaskAppService service = new KnowledgeDocumentImportTaskAppService();
|
||||
Method method = KnowledgeDocumentImportTaskAppService.class.getDeclaredMethod(
|
||||
"resolveParseFailureCode",
|
||||
Throwable.class
|
||||
);
|
||||
method.setAccessible(true);
|
||||
|
||||
Object code = method.invoke(service,
|
||||
new RuntimeException("wrapped", new SocketTimeoutException("timeout")));
|
||||
|
||||
Assert.assertEquals("parse_service_timeout", code);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证文档源读取失败优先于底层 UnknownHostException 分类，避免误报 MinerU 不可用。
|
||||
*
|
||||
* @throws Exception 反射调用异常
|
||||
*/
|
||||
@Test
|
||||
public void resolveParseFailureCodeShouldPreferDocumentSourceFailure() throws Exception {
|
||||
KnowledgeDocumentImportTaskAppService service = new KnowledgeDocumentImportTaskAppService();
|
||||
Method method = KnowledgeDocumentImportTaskAppService.class.getDeclaredMethod(
|
||||
"resolveParseFailureCode",
|
||||
Throwable.class
|
||||
);
|
||||
method.setAccessible(true);
|
||||
|
||||
DocumentParseBridgeException sourceError = DocumentParseBridgeException.sourceLoadFailed(
|
||||
"下载文档 URL 失败",
|
||||
new UnknownHostException("远端文档地址不允许访问非公网目标")
|
||||
);
|
||||
Object code = method.invoke(service, new RuntimeException("wrapped", sourceError));
|
||||
|
||||
Assert.assertEquals("document_source_unavailable", code);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证解析桥接层的确定性输入错误会转换为不可重试错误码。
|
||||
*
|
||||
* @throws Exception 反射调用异常
|
||||
*/
|
||||
@Test
|
||||
public void resolveParseFailureCodeShouldClassifyPermanentBridgeFailures()
|
||||
throws Exception {
|
||||
KnowledgeDocumentImportTaskAppService service =
|
||||
new KnowledgeDocumentImportTaskAppService();
|
||||
Method method = KnowledgeDocumentImportTaskAppService.class
|
||||
.getDeclaredMethod("resolveParseFailureCode", Throwable.class);
|
||||
method.setAccessible(true);
|
||||
|
||||
Object unsupported = method.invoke(
|
||||
service,
|
||||
DocumentParseBridgeException.unsupportedSource("不支持的文件类型")
|
||||
);
|
||||
Object invalidRequest = method.invoke(
|
||||
service,
|
||||
DocumentParseBridgeException.requestBuildFailed("解析请求无效")
|
||||
);
|
||||
|
||||
Assert.assertEquals("unsupported_document_source", unsupported);
|
||||
Assert.assertEquals("invalid_parse_request", invalidRequest);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证未识别的系统异常仍会获得稳定错误码并允许后续人工重试。
|
||||
*
|
||||
* @throws Exception 反射调用异常
|
||||
*/
|
||||
@Test
|
||||
public void resolveParseFailureCodeShouldStabilizeUnknownFailures()
|
||||
throws Exception {
|
||||
KnowledgeDocumentImportTaskAppService service =
|
||||
new KnowledgeDocumentImportTaskAppService();
|
||||
Method method = KnowledgeDocumentImportTaskAppService.class
|
||||
.getDeclaredMethod("resolveParseFailureCode", Throwable.class);
|
||||
method.setAccessible(true);
|
||||
|
||||
Object code = method.invoke(service, new IllegalStateException("代码异常"));
|
||||
|
||||
Assert.assertEquals("parse_failed", code);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证未知系统异常和临时源读取失败允许人工重试，确定性输入错误除外。
|
||||
*
|
||||
* @throws Exception 反射调用异常
|
||||
*/
|
||||
@Test
|
||||
public void parseRetryabilityShouldDefaultToRecoverable()
|
||||
throws Exception {
|
||||
KnowledgeDocumentImportTaskAppService service =
|
||||
new KnowledgeDocumentImportTaskAppService();
|
||||
Method method = KnowledgeDocumentImportTaskAppService.class
|
||||
.getDeclaredMethod("isRetryableParseFailure", String.class);
|
||||
method.setAccessible(true);
|
||||
|
||||
Assert.assertTrue((Boolean) method.invoke(service, new Object[] {null}));
|
||||
Assert.assertTrue((Boolean) method.invoke(
|
||||
service, "document_source_unavailable"));
|
||||
Assert.assertTrue((Boolean) method.invoke(service, "parse_failed"));
|
||||
Assert.assertFalse((Boolean) method.invoke(
|
||||
service, "unsupported_document_source"));
|
||||
Assert.assertFalse((Boolean) method.invoke(
|
||||
service, "invalid_parse_request"));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证批次上传文件通过受信任存储服务读取为字节，避免内部附件 URL 进入公网 URL 校验。
|
||||
*
|
||||
* @throws Exception 反射调用或存储桩异常
|
||||
*/
|
||||
@Test
|
||||
public void buildBridgeSourceRefShouldReadBatchFileFromStorage() throws Exception {
|
||||
byte[] content = "batch-document".getBytes(StandardCharsets.UTF_8);
|
||||
String storedUrl = "http://127.0.0.1:39000/easyflow/attachment/test.docx";
|
||||
FileStorageService storageService = Mockito.mock(FileStorageService.class);
|
||||
Mockito.when(storageService.readStream(storedUrl))
|
||||
.thenReturn(new ByteArrayInputStream(content));
|
||||
|
||||
DocumentImportBulkProperties properties = new DocumentImportBulkProperties();
|
||||
KnowledgeDocumentImportTaskAppService service = new KnowledgeDocumentImportTaskAppService();
|
||||
setField(service, "storageService", storageService);
|
||||
setField(service, "bulkProperties", properties);
|
||||
|
||||
DocumentImportTask task = new DocumentImportTask();
|
||||
task.setBatchId(BigInteger.valueOf(61));
|
||||
tech.easyflow.ai.entity.Document document = new tech.easyflow.ai.entity.Document();
|
||||
document.setTitle("test.docx");
|
||||
document.setDocumentPath(storedUrl);
|
||||
|
||||
Method method = KnowledgeDocumentImportTaskAppService.class.getDeclaredMethod(
|
||||
"buildBridgeSourceRef",
|
||||
DocumentImportTask.class,
|
||||
tech.easyflow.ai.entity.Document.class,
|
||||
String.class
|
||||
);
|
||||
method.setAccessible(true);
|
||||
DocumentSourceRef sourceRef = (DocumentSourceRef) method.invoke(service, task, document, "docx");
|
||||
|
||||
Assert.assertArrayEquals(content, sourceRef.getContentBytes());
|
||||
Assert.assertNull(sourceRef.getFilePath());
|
||||
Assert.assertEquals(Long.valueOf(content.length), sourceRef.getSize());
|
||||
Mockito.verify(storageService).readStream(storedUrl);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证非批次远程文档仍保留 URL，由文档源加载器继续执行 SSRF 防护。
|
||||
*
|
||||
* @throws Exception 反射调用异常
|
||||
*/
|
||||
@Test
|
||||
public void buildBridgeSourceRefShouldKeepExternalUrlForSsrfValidation() throws Exception {
|
||||
FileStorageService storageService = Mockito.mock(FileStorageService.class);
|
||||
KnowledgeDocumentImportTaskAppService service = new KnowledgeDocumentImportTaskAppService();
|
||||
setField(service, "storageService", storageService);
|
||||
|
||||
DocumentImportTask task = new DocumentImportTask();
|
||||
tech.easyflow.ai.entity.Document document = new tech.easyflow.ai.entity.Document();
|
||||
document.setTitle("external.pdf");
|
||||
document.setDocumentPath("https://example.com/external.pdf");
|
||||
|
||||
Method method = KnowledgeDocumentImportTaskAppService.class.getDeclaredMethod(
|
||||
"buildBridgeSourceRef",
|
||||
DocumentImportTask.class,
|
||||
tech.easyflow.ai.entity.Document.class,
|
||||
String.class
|
||||
);
|
||||
method.setAccessible(true);
|
||||
DocumentSourceRef sourceRef = (DocumentSourceRef) method.invoke(service, task, document, "pdf");
|
||||
|
||||
Assert.assertEquals("https://example.com/external.pdf", sourceRef.getFilePath());
|
||||
Assert.assertNull(sourceRef.getContentBytes());
|
||||
Mockito.verifyNoInteractions(storageService);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证向量化失败会按整文档失败语义重置进度，并刷新任务错误信息。
|
||||
*
|
||||
@@ -59,11 +744,16 @@ public class KnowledgeDocumentImportTaskAppServiceTest {
|
||||
persistedDocument.setLastTaskError("旧错误");
|
||||
|
||||
AtomicReference<tech.easyflow.ai.entity.Document> updatedDocumentRef = new AtomicReference<tech.easyflow.ai.entity.Document>();
|
||||
AtomicReference<DocumentImportTask> updatedTaskRef = new AtomicReference<DocumentImportTask>();
|
||||
DocumentImportTaskMapper taskMapper = Mockito.mock(DocumentImportTaskMapper.class);
|
||||
Mockito.when(taskMapper.finishOwned(
|
||||
Mockito.any(), Mockito.anyString(), Mockito.anyString(),
|
||||
Mockito.any(), Mockito.anyString(), Mockito.any(),
|
||||
Mockito.any()
|
||||
)).thenReturn(1);
|
||||
|
||||
KnowledgeDocumentImportTaskAppService service = new KnowledgeDocumentImportTaskAppService();
|
||||
setField(service, "documentMapper", mockDocumentMapper(persistedDocument, updatedDocumentRef));
|
||||
setField(service, "documentImportTaskService", mockDocumentImportTaskService(updatedTaskRef));
|
||||
setField(service, "documentImportTaskMapper", taskMapper);
|
||||
setField(service, "documentImportTaskStatusStreamService", new NoopTaskStatusStreamService());
|
||||
|
||||
DocumentImportTask task = new DocumentImportTask();
|
||||
@@ -72,6 +762,7 @@ public class KnowledgeDocumentImportTaskAppServiceTest {
|
||||
task.setKnowledgeId(knowledgeId);
|
||||
task.setStatus(DocumentImportTaskStatus.RUNNING.name());
|
||||
task.setErrorSummary("旧错误");
|
||||
task.setExecutionToken("attempt-token");
|
||||
|
||||
tech.easyflow.ai.entity.Document inputDocument = new tech.easyflow.ai.entity.Document();
|
||||
inputDocument.setId(documentId);
|
||||
@@ -93,11 +784,62 @@ public class KnowledgeDocumentImportTaskAppServiceTest {
|
||||
Assert.assertEquals(Integer.valueOf(8), updatedDocument.getFailedChunks());
|
||||
Assert.assertEquals(Integer.valueOf(0), updatedDocument.getProgressPercent());
|
||||
Assert.assertEquals("新错误", updatedDocument.getLastTaskError());
|
||||
Assert.assertEquals("index_failed",
|
||||
updatedDocument.getOptions().get(DocumentImportKeys.KEY_DOCUMENT_TASK_ERROR_CODE));
|
||||
|
||||
DocumentImportTask updatedTask = updatedTaskRef.get();
|
||||
Assert.assertNotNull(updatedTask);
|
||||
Assert.assertEquals(DocumentImportTaskStatus.FAILED.name(), updatedTask.getStatus());
|
||||
Assert.assertEquals("新错误", updatedTask.getErrorSummary());
|
||||
Mockito.verify(taskMapper).finishOwned(
|
||||
Mockito.eq(task.getId()),
|
||||
Mockito.eq("attempt-token"),
|
||||
Mockito.eq(DocumentImportTaskStatus.FAILED.name()),
|
||||
Mockito.eq("新错误"),
|
||||
Mockito.eq("index_failed"),
|
||||
Mockito.any(Date.class),
|
||||
Mockito.any()
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证看门狗已经收口任务后，迟到执行者不能覆盖文档终态。
|
||||
*
|
||||
* @throws Exception 反射调用异常
|
||||
*/
|
||||
@Test
|
||||
public void markIndexFailedShouldIgnoreLostTaskOwnership() throws Exception {
|
||||
BigInteger documentId = BigInteger.valueOf(81);
|
||||
tech.easyflow.ai.entity.Document persistedDocument = new tech.easyflow.ai.entity.Document();
|
||||
persistedDocument.setId(documentId);
|
||||
persistedDocument.setProcessStatus(DocumentProcessStatus.PARSE_FAILED.name());
|
||||
AtomicReference<tech.easyflow.ai.entity.Document> updatedDocumentRef =
|
||||
new AtomicReference<tech.easyflow.ai.entity.Document>();
|
||||
DocumentImportTaskMapper taskMapper = Mockito.mock(DocumentImportTaskMapper.class);
|
||||
Mockito.when(taskMapper.updateByQuery(
|
||||
Mockito.any(DocumentImportTask.class),
|
||||
Mockito.any()
|
||||
)).thenReturn(0);
|
||||
|
||||
KnowledgeDocumentImportTaskAppService service = new KnowledgeDocumentImportTaskAppService();
|
||||
setField(service, "documentMapper", mockDocumentMapper(persistedDocument, updatedDocumentRef));
|
||||
setField(service, "documentImportTaskMapper", taskMapper);
|
||||
|
||||
DocumentImportTask task = new DocumentImportTask();
|
||||
task.setId(BigInteger.valueOf(82));
|
||||
task.setDocumentId(documentId);
|
||||
task.setStatus(DocumentImportTaskStatus.RUNNING.name());
|
||||
tech.easyflow.ai.entity.Document inputDocument = new tech.easyflow.ai.entity.Document();
|
||||
inputDocument.setId(documentId);
|
||||
|
||||
Method method = KnowledgeDocumentImportTaskAppService.class.getDeclaredMethod(
|
||||
"markIndexFailed",
|
||||
DocumentImportTask.class,
|
||||
tech.easyflow.ai.entity.Document.class,
|
||||
String.class
|
||||
);
|
||||
method.setAccessible(true);
|
||||
method.invoke(service, task, inputDocument, "迟到错误");
|
||||
|
||||
Assert.assertNull(updatedDocumentRef.get());
|
||||
Assert.assertEquals(DocumentProcessStatus.PARSE_FAILED.name(),
|
||||
persistedDocument.getProcessStatus());
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -232,6 +974,8 @@ public class KnowledgeDocumentImportTaskAppServiceTest {
|
||||
|
||||
Assert.assertEquals(2, chunks.size());
|
||||
DocumentChunk firstChunk = chunks.get(0);
|
||||
Assert.assertNotNull(firstChunk.getId());
|
||||
Assert.assertNotEquals(firstChunk.getId(), chunks.get(1).getId());
|
||||
Assert.assertTrue(firstChunk.getContent().contains("Slide 1"));
|
||||
Assert.assertTrue(firstChunk.getContent().contains("本页介绍季度目标"));
|
||||
Assert.assertEquals("https://example.com/slides/slide-001.png",
|
||||
@@ -352,6 +1096,61 @@ public class KnowledgeDocumentImportTaskAppServiceTest {
|
||||
Assert.assertTrue(chunks.isEmpty());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证向量存储回滚会去重分块 ID，并在 Milvus 删除失败后停止清理搜索索引。
|
||||
*
|
||||
* @throws Exception 反射调用异常
|
||||
*/
|
||||
@Test
|
||||
public void rollbackStoredChunksShouldValidateVectorDeleteResult() throws Exception {
|
||||
KnowledgeDocumentImportTaskAppService service = new KnowledgeDocumentImportTaskAppService();
|
||||
DocumentStore documentStore = Mockito.mock(DocumentStore.class);
|
||||
DocumentSearcher searcher = Mockito.mock(DocumentSearcher.class);
|
||||
StoreOptions storeOptions = StoreOptions.ofCollectionName("knowledge-test");
|
||||
Mockito.when(documentStore.delete(Mockito.anyCollection(), Mockito.same(storeOptions)))
|
||||
.thenReturn(StoreResult.fail("milvus unavailable"));
|
||||
|
||||
DocumentCollection knowledge = new DocumentCollection();
|
||||
knowledge.setId(BigInteger.valueOf(501));
|
||||
Class<?> contextClass = Class.forName(
|
||||
KnowledgeDocumentImportTaskAppService.class.getName() + "$StoreExecutionContext"
|
||||
);
|
||||
Constructor<?> constructor = contextClass.getDeclaredConstructor(
|
||||
DocumentCollection.class,
|
||||
EmbeddingModel.class,
|
||||
DocumentStore.class,
|
||||
StoreOptions.class,
|
||||
DocumentSearcher.class
|
||||
);
|
||||
constructor.setAccessible(true);
|
||||
Object context = constructor.newInstance(knowledge, null, documentStore, storeOptions, searcher);
|
||||
|
||||
DocumentChunk first = new DocumentChunk();
|
||||
first.setId(BigInteger.valueOf(601));
|
||||
DocumentChunk duplicate = new DocumentChunk();
|
||||
duplicate.setId(first.getId());
|
||||
Method method = KnowledgeDocumentImportTaskAppService.class.getDeclaredMethod(
|
||||
"rollbackStoredChunks",
|
||||
BigInteger.class,
|
||||
BigInteger.class,
|
||||
contextClass,
|
||||
List.class
|
||||
);
|
||||
method.setAccessible(true);
|
||||
method.invoke(service,
|
||||
BigInteger.valueOf(701),
|
||||
BigInteger.valueOf(702),
|
||||
context,
|
||||
List.of(first, duplicate));
|
||||
|
||||
@SuppressWarnings("rawtypes")
|
||||
ArgumentCaptor<Collection> idsCaptor = ArgumentCaptor.forClass(Collection.class);
|
||||
Mockito.verify(documentStore).delete(idsCaptor.capture(), Mockito.same(storeOptions));
|
||||
Assert.assertEquals(1, idsCaptor.getValue().size());
|
||||
Assert.assertTrue(idsCaptor.getValue().contains(first.getId()));
|
||||
Mockito.verifyNoInteractions(searcher);
|
||||
}
|
||||
|
||||
private static DocumentMapper mockDocumentMapper(tech.easyflow.ai.entity.Document persistedDocument,
|
||||
AtomicReference<tech.easyflow.ai.entity.Document> updatedDocumentRef) {
|
||||
return (DocumentMapper) Proxy.newProxyInstance(
|
||||
@@ -370,20 +1169,6 @@ public class KnowledgeDocumentImportTaskAppServiceTest {
|
||||
);
|
||||
}
|
||||
|
||||
private static DocumentImportTaskService mockDocumentImportTaskService(AtomicReference<DocumentImportTask> updatedTaskRef) {
|
||||
return (DocumentImportTaskService) Proxy.newProxyInstance(
|
||||
DocumentImportTaskService.class.getClassLoader(),
|
||||
new Class<?>[]{DocumentImportTaskService.class},
|
||||
(proxy, method, args) -> {
|
||||
if ("updateById".equals(method.getName())) {
|
||||
updatedTaskRef.set((DocumentImportTask) args[0]);
|
||||
return true;
|
||||
}
|
||||
return defaultValue(method.getReturnType());
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
private static FileStorageService mockFileStorageService(AtomicReference<String> savedPrePathRef,
|
||||
AtomicReference<String> savedFilenameRef) {
|
||||
return (FileStorageService) Proxy.newProxyInstance(
|
||||
|
||||
@@ -0,0 +1,431 @@
|
||||
package tech.easyflow.ai.documentimport.task;
|
||||
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
import org.mockito.ArgumentCaptor;
|
||||
import org.mockito.Mockito;
|
||||
import org.springframework.web.multipart.MultipartFile;
|
||||
import tech.easyflow.ai.documentimport.DocumentImportBatchCreateContext;
|
||||
import tech.easyflow.ai.documentimport.DocumentImportBatchDtos;
|
||||
import tech.easyflow.ai.documentimport.ImportCallerContext;
|
||||
import tech.easyflow.ai.documentimport.ImportCallerType;
|
||||
import tech.easyflow.ai.documentimport.PublicDocumentImportDtos;
|
||||
import tech.easyflow.ai.entity.DocumentImportBatch;
|
||||
import tech.easyflow.ai.enums.DocumentImportBatchStatus;
|
||||
import tech.easyflow.ai.mapper.DocumentImportBatchMapper;
|
||||
import tech.easyflow.ai.service.DocumentImportBatchItemService;
|
||||
import tech.easyflow.ai.service.DocumentImportBatchService;
|
||||
import tech.easyflow.common.web.exceptions.BusinessException;
|
||||
|
||||
import java.io.ByteArrayInputStream;
|
||||
import java.math.BigInteger;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Date;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* Public API 批量导入门面边界与服务端去重测试。
|
||||
*
|
||||
* @author Codex
|
||||
* @since 2026-08-02
|
||||
*/
|
||||
public class KnowledgeImportBatchFacadeTest {
|
||||
|
||||
/**
|
||||
* 验证单次文件数超过 200 时在创建批次前拒绝。
|
||||
*/
|
||||
@Test
|
||||
public void submitShouldRejectMoreThanTwoHundredFiles() {
|
||||
TestContext context = createContext();
|
||||
PublicDocumentImportDtos.BatchMetadata metadata =
|
||||
new PublicDocumentImportDtos.BatchMetadata();
|
||||
metadata.setKnowledgeId(BigInteger.ONE);
|
||||
List<DocumentImportBatchDtos.ManifestItem> manifest = new ArrayList<>();
|
||||
List<MultipartFile> files = new ArrayList<>();
|
||||
for (int index = 0; index < 201; index++) {
|
||||
String name = "file-" + index + ".txt";
|
||||
manifest.add(manifest(name, name, 1L));
|
||||
files.add(file(name, new byte[]{'a'}));
|
||||
}
|
||||
metadata.setFiles(manifest);
|
||||
|
||||
try {
|
||||
context.facade.submit(context.caller, metadata, files);
|
||||
Assert.fail("Expected file count rejection");
|
||||
} catch (BusinessException expected) {
|
||||
Assert.assertTrue(expected.getMessage().contains("最多上传200个文件"));
|
||||
}
|
||||
Mockito.verifyNoInteractions(context.batchAppService);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证文件夹清单拒绝父目录穿越。
|
||||
*/
|
||||
@Test
|
||||
public void submitShouldRejectRelativePathTraversal() {
|
||||
TestContext context = createContext();
|
||||
PublicDocumentImportDtos.BatchMetadata metadata =
|
||||
metadata("demo.txt", "../demo.txt", 4L);
|
||||
|
||||
try {
|
||||
context.facade.submit(
|
||||
context.caller,
|
||||
metadata,
|
||||
List.of(file(
|
||||
"demo.txt",
|
||||
"demo".getBytes(StandardCharsets.UTF_8)
|
||||
))
|
||||
);
|
||||
Assert.fail("Expected relative path rejection");
|
||||
} catch (BusinessException expected) {
|
||||
Assert.assertTrue(expected.getMessage().contains("相对路径无效"));
|
||||
}
|
||||
Mockito.verifyNoInteractions(context.batchAppService);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证 Public API 可省略 fileSize，内部批次使用服务端实测大小。
|
||||
*/
|
||||
@Test
|
||||
public void submitShouldUseMeasuredSizeWhenMetadataOmitsFileSize() {
|
||||
TestContext context = createContext();
|
||||
PublicDocumentImportDtos.BatchMetadata metadata =
|
||||
metadata("demo.txt", "folder/demo.txt", 4L);
|
||||
metadata.getFiles().get(0).setFileSize(null);
|
||||
byte[] content = "demo".getBytes(StandardCharsets.UTF_8);
|
||||
|
||||
DocumentImportBatchDtos.CreateResponse created =
|
||||
new DocumentImportBatchDtos.CreateResponse();
|
||||
created.setBatchId(BigInteger.valueOf(91));
|
||||
DocumentImportBatchDtos.ItemResponse createdItem =
|
||||
new DocumentImportBatchDtos.ItemResponse();
|
||||
createdItem.setItemId(BigInteger.valueOf(92));
|
||||
created.setItems(List.of(createdItem));
|
||||
Mockito.when(context.batchAppService.createBatch(
|
||||
Mockito.any(), Mockito.any()
|
||||
)).thenReturn(created);
|
||||
Mockito.when(context.itemService.updateById(Mockito.any()))
|
||||
.thenReturn(true);
|
||||
DocumentImportBatch batch = new DocumentImportBatch();
|
||||
batch.setId(created.getBatchId());
|
||||
batch.setKnowledgeId(BigInteger.ONE);
|
||||
batch.setStatus(DocumentImportBatchStatus.RUNNING.name());
|
||||
batch.setTotalCount(1);
|
||||
batch.setTotalBytes((long) content.length);
|
||||
batch.setCreated(new Date());
|
||||
Mockito.when(context.batchAppService.requireOwnedBatch(
|
||||
BigInteger.ONE, created.getBatchId(), context.caller
|
||||
)).thenReturn(batch);
|
||||
|
||||
context.facade.submit(
|
||||
context.caller,
|
||||
metadata,
|
||||
List.of(file("demo.txt", content))
|
||||
);
|
||||
|
||||
ArgumentCaptor<DocumentImportBatchDtos.CreateRequest> request =
|
||||
ArgumentCaptor.forClass(DocumentImportBatchDtos.CreateRequest.class);
|
||||
Mockito.verify(context.batchAppService).createBatch(
|
||||
request.capture(), Mockito.any());
|
||||
Assert.assertEquals(
|
||||
Long.valueOf(content.length),
|
||||
request.getValue().getFiles().get(0).getFileSize()
|
||||
);
|
||||
Assert.assertNull(metadata.getFiles().get(0).getFileSize());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证相同提交由服务端指纹复用，包含仍在上传的并发请求。
|
||||
*/
|
||||
@Test
|
||||
public void submitShouldReuseTaskByServerFingerprint() {
|
||||
TestContext context = createContext();
|
||||
PublicDocumentImportDtos.BatchMetadata metadata =
|
||||
metadata("demo.txt", "folder/demo.txt", 4L);
|
||||
List<MultipartFile> files = List.of(file(
|
||||
"demo.txt",
|
||||
"demo".getBytes(StandardCharsets.UTF_8)
|
||||
));
|
||||
|
||||
DocumentImportBatchDtos.CreateResponse created =
|
||||
new DocumentImportBatchDtos.CreateResponse();
|
||||
created.setBatchId(BigInteger.valueOf(91));
|
||||
DocumentImportBatchDtos.ItemResponse createdItem =
|
||||
new DocumentImportBatchDtos.ItemResponse();
|
||||
createdItem.setItemId(BigInteger.valueOf(92));
|
||||
created.setItems(List.of(createdItem));
|
||||
Mockito.when(context.batchAppService.createBatch(
|
||||
Mockito.any(), Mockito.any()
|
||||
)).thenReturn(created);
|
||||
Mockito.when(context.itemService.updateById(Mockito.any()))
|
||||
.thenReturn(true);
|
||||
DocumentImportBatch batch = new DocumentImportBatch();
|
||||
batch.setId(created.getBatchId());
|
||||
batch.setKnowledgeId(BigInteger.ONE);
|
||||
batch.setStatus(DocumentImportBatchStatus.RUNNING.name());
|
||||
batch.setTotalCount(1);
|
||||
batch.setTotalBytes(4L);
|
||||
batch.setCreated(new Date());
|
||||
Mockito.when(context.batchAppService.requireOwnedBatch(
|
||||
BigInteger.ONE, created.getBatchId(), context.caller
|
||||
)).thenReturn(batch);
|
||||
|
||||
context.facade.submit(context.caller, metadata, files);
|
||||
|
||||
ArgumentCaptor<tech.easyflow.ai.entity.DocumentImportBatchItem> hashUpdate =
|
||||
ArgumentCaptor.forClass(tech.easyflow.ai.entity.DocumentImportBatchItem.class);
|
||||
Mockito.verify(context.itemService).updateById(hashUpdate.capture());
|
||||
Assert.assertNotNull(hashUpdate.getValue().getContentSha256());
|
||||
Assert.assertEquals(createdItem.getItemId(), hashUpdate.getValue().getId());
|
||||
ArgumentCaptor<DocumentImportBatchCreateContext> createContext =
|
||||
ArgumentCaptor.forClass(DocumentImportBatchCreateContext.class);
|
||||
Mockito.verify(context.batchAppService).createBatch(
|
||||
Mockito.any(), createContext.capture());
|
||||
batch.setRequestDigest(createContext.getValue().getRequestDigest());
|
||||
Mockito.when(context.batchService.getOne(
|
||||
Mockito.any(com.mybatisflex.core.query.QueryWrapper.class)))
|
||||
.thenReturn(batch);
|
||||
|
||||
PublicDocumentImportDtos.SubmitResponse repeated =
|
||||
context.facade.submit(
|
||||
context.caller, metadata, files);
|
||||
|
||||
Assert.assertEquals(batch.getId(), repeated.getTaskId());
|
||||
batch.setStatus(DocumentImportBatchStatus.UPLOADING.name());
|
||||
PublicDocumentImportDtos.SubmitResponse uploading =
|
||||
context.facade.submit(context.caller, metadata, files);
|
||||
|
||||
Assert.assertEquals(batch.getId(), uploading.getTaskId());
|
||||
Mockito.verify(context.batchAppService, Mockito.times(1))
|
||||
.createBatch(Mockito.any(), Mockito.any());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证长时间无进展的上传任务会先取消并释放指纹，再创建新任务。
|
||||
*/
|
||||
@Test
|
||||
public void submitShouldReplaceStaleIncompleteTask() {
|
||||
TestContext context = createContext();
|
||||
PublicDocumentImportDtos.BatchMetadata metadata =
|
||||
metadata("demo.txt", "folder/demo.txt", 4L);
|
||||
List<MultipartFile> files = List.of(file(
|
||||
"demo.txt",
|
||||
"demo".getBytes(StandardCharsets.UTF_8)
|
||||
));
|
||||
DocumentImportBatch stale = new DocumentImportBatch();
|
||||
stale.setId(BigInteger.valueOf(81));
|
||||
stale.setKnowledgeId(BigInteger.ONE);
|
||||
stale.setStatus(DocumentImportBatchStatus.UPLOADING.name());
|
||||
stale.setModified(new Date(
|
||||
System.currentTimeMillis() - 31L * 60L * 1000L
|
||||
));
|
||||
Mockito.when(context.batchService.getOne(
|
||||
Mockito.any(com.mybatisflex.core.query.QueryWrapper.class)))
|
||||
.thenReturn(stale)
|
||||
.thenReturn(null);
|
||||
Mockito.when(context.batchAppService.cancelStaleBatch(
|
||||
Mockito.eq(BigInteger.ONE),
|
||||
Mockito.eq(stale.getId()),
|
||||
Mockito.eq(context.caller),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(true);
|
||||
Mockito.when(context.batchMapper.releaseSubmissionFingerprint(
|
||||
Mockito.eq(stale.getId()),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class),
|
||||
Mockito.any(Date.class)
|
||||
)).thenReturn(1);
|
||||
|
||||
DocumentImportBatchDtos.CreateResponse created =
|
||||
new DocumentImportBatchDtos.CreateResponse();
|
||||
created.setBatchId(BigInteger.valueOf(91));
|
||||
DocumentImportBatchDtos.ItemResponse item =
|
||||
new DocumentImportBatchDtos.ItemResponse();
|
||||
item.setItemId(BigInteger.valueOf(92));
|
||||
created.setItems(List.of(item));
|
||||
Mockito.when(context.batchAppService.createBatch(
|
||||
Mockito.any(), Mockito.any()
|
||||
)).thenReturn(created);
|
||||
Mockito.when(context.itemService.updateById(Mockito.any()))
|
||||
.thenReturn(true);
|
||||
DocumentImportBatch replacement = new DocumentImportBatch();
|
||||
replacement.setId(created.getBatchId());
|
||||
replacement.setKnowledgeId(BigInteger.ONE);
|
||||
replacement.setStatus(DocumentImportBatchStatus.RUNNING.name());
|
||||
replacement.setTotalCount(1);
|
||||
replacement.setTotalBytes(4L);
|
||||
replacement.setCreated(new Date());
|
||||
Mockito.when(context.batchAppService.requireOwnedBatch(
|
||||
BigInteger.ONE, created.getBatchId(), context.caller
|
||||
)).thenReturn(replacement);
|
||||
|
||||
PublicDocumentImportDtos.SubmitResponse response =
|
||||
context.facade.submit(context.caller, metadata, files);
|
||||
|
||||
Assert.assertEquals(replacement.getId(), response.getTaskId());
|
||||
Mockito.verify(context.batchAppService).cancelStaleBatch(
|
||||
Mockito.eq(BigInteger.ONE),
|
||||
Mockito.eq(stale.getId()),
|
||||
Mockito.eq(context.caller),
|
||||
Mockito.any(Date.class)
|
||||
);
|
||||
Mockito.verify(context.batchMapper).releaseSubmissionFingerprint(
|
||||
Mockito.eq(stale.getId()),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class),
|
||||
Mockito.any(Date.class)
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证显式空文件键数组会被拒绝，省略时由服务端重试全部异常文件。
|
||||
*/
|
||||
@Test
|
||||
public void retryShouldDistinguishOmittedAndEmptyFileKeys() {
|
||||
TestContext context = createContext();
|
||||
PublicDocumentImportDtos.RetryRequest empty =
|
||||
new PublicDocumentImportDtos.RetryRequest();
|
||||
empty.setTaskId(BigInteger.ONE);
|
||||
empty.setFileKeys(List.of());
|
||||
|
||||
try {
|
||||
context.facade.retry(context.caller, empty);
|
||||
Assert.fail("Expected empty fileKeys rejection");
|
||||
} catch (BusinessException expected) {
|
||||
Assert.assertTrue(expected.getMessage().contains("不能为空数组"));
|
||||
}
|
||||
|
||||
PublicDocumentImportDtos.RetryRequest omitted =
|
||||
new PublicDocumentImportDtos.RetryRequest();
|
||||
omitted.setTaskId(BigInteger.ONE);
|
||||
tech.easyflow.ai.documentimport.DocumentImportBatchRetryResult result =
|
||||
new tech.easyflow.ai.documentimport.DocumentImportBatchRetryResult(
|
||||
BigInteger.ONE,
|
||||
DocumentImportBatchStatus.RUNNING.name(),
|
||||
1,
|
||||
2
|
||||
);
|
||||
Mockito.when(context.batchAppService.retryOwnedBatch(
|
||||
BigInteger.ONE,
|
||||
context.caller,
|
||||
java.util.Set.of()
|
||||
)).thenReturn(result);
|
||||
|
||||
PublicDocumentImportDtos.RetryResponse response =
|
||||
context.facade.retry(context.caller, omitted);
|
||||
|
||||
Assert.assertEquals(Integer.valueOf(2), response.getRetriedCount());
|
||||
}
|
||||
|
||||
/**
|
||||
* 创建单文件元数据。
|
||||
*
|
||||
* @param fileName 文件名
|
||||
* @param relativePath 相对路径
|
||||
* @param fileSize 文件大小
|
||||
* @return 元数据
|
||||
*/
|
||||
private PublicDocumentImportDtos.BatchMetadata metadata(
|
||||
String fileName,
|
||||
String relativePath,
|
||||
long fileSize) {
|
||||
PublicDocumentImportDtos.BatchMetadata metadata =
|
||||
new PublicDocumentImportDtos.BatchMetadata();
|
||||
metadata.setKnowledgeId(BigInteger.ONE);
|
||||
metadata.setFiles(List.of(
|
||||
manifest(fileName, relativePath, fileSize)));
|
||||
return metadata;
|
||||
}
|
||||
|
||||
/**
|
||||
* 创建清单项。
|
||||
*
|
||||
* @param fileName 文件名
|
||||
* @param relativePath 相对路径
|
||||
* @param fileSize 文件大小
|
||||
* @return 清单项
|
||||
*/
|
||||
private DocumentImportBatchDtos.ManifestItem manifest(
|
||||
String fileName,
|
||||
String relativePath,
|
||||
long fileSize) {
|
||||
DocumentImportBatchDtos.ManifestItem item =
|
||||
new DocumentImportBatchDtos.ManifestItem();
|
||||
item.setClientFileKey(relativePath);
|
||||
item.setFileName(fileName);
|
||||
item.setRelativePath(relativePath);
|
||||
item.setFileSize(fileSize);
|
||||
return item;
|
||||
}
|
||||
|
||||
/**
|
||||
* 创建可重复读取的 MultipartFile。
|
||||
*
|
||||
* @param fileName 文件名
|
||||
* @param bytes 文件内容
|
||||
* @return Multipart 文件
|
||||
*/
|
||||
private MultipartFile file(String fileName, byte[] bytes) {
|
||||
MultipartFile file = Mockito.mock(MultipartFile.class);
|
||||
try {
|
||||
Mockito.when(file.getInputStream())
|
||||
.thenAnswer(invocation -> new ByteArrayInputStream(bytes));
|
||||
} catch (java.io.IOException error) {
|
||||
throw new IllegalStateException(error);
|
||||
}
|
||||
Mockito.when(file.getOriginalFilename()).thenReturn(fileName);
|
||||
Mockito.when(file.getSize()).thenReturn((long) bytes.length);
|
||||
Mockito.when(file.isEmpty()).thenReturn(bytes.length == 0);
|
||||
return file;
|
||||
}
|
||||
|
||||
/**
|
||||
* 创建测试依赖。
|
||||
*
|
||||
* @return 测试上下文
|
||||
*/
|
||||
private TestContext createContext() {
|
||||
DocumentImportBatchAppService batchAppService =
|
||||
Mockito.mock(DocumentImportBatchAppService.class);
|
||||
DocumentImportBatchTracker batchTracker =
|
||||
Mockito.mock(DocumentImportBatchTracker.class);
|
||||
DocumentImportBatchService batchService =
|
||||
Mockito.mock(DocumentImportBatchService.class);
|
||||
DocumentImportBatchItemService itemService =
|
||||
Mockito.mock(DocumentImportBatchItemService.class);
|
||||
DocumentImportBatchMapper batchMapper =
|
||||
Mockito.mock(DocumentImportBatchMapper.class);
|
||||
KnowledgeImportBatchFacade facade = new KnowledgeImportBatchFacade(
|
||||
batchAppService,
|
||||
batchTracker,
|
||||
batchService,
|
||||
itemService,
|
||||
batchMapper
|
||||
);
|
||||
return new TestContext(
|
||||
facade,
|
||||
batchAppService,
|
||||
batchService,
|
||||
itemService,
|
||||
batchMapper,
|
||||
new ImportCallerContext(
|
||||
ImportCallerType.PUBLIC_API,
|
||||
BigInteger.valueOf(77)
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 测试依赖集合。
|
||||
*/
|
||||
private record TestContext(
|
||||
KnowledgeImportBatchFacade facade,
|
||||
DocumentImportBatchAppService batchAppService,
|
||||
DocumentImportBatchService batchService,
|
||||
DocumentImportBatchItemService itemService,
|
||||
DocumentImportBatchMapper batchMapper,
|
||||
ImportCallerContext caller
|
||||
) {
|
||||
}
|
||||
}
|
||||
@@ -10,6 +10,7 @@ import tech.easyflow.ai.config.SearcherFactory;
|
||||
import tech.easyflow.ai.enums.DocumentProcessStatus;
|
||||
import tech.easyflow.ai.mapper.DocumentChunkMapper;
|
||||
import tech.easyflow.ai.mapper.DocumentMapper;
|
||||
import tech.easyflow.ai.mapper.FaqItemMapper;
|
||||
|
||||
import java.io.Serializable;
|
||||
import java.lang.reflect.Field;
|
||||
@@ -92,6 +93,89 @@ public class DocumentCollectionServiceImplTest {
|
||||
Assert.assertEquals(completedChunkId, result.get(0).getId());
|
||||
Assert.assertEquals("completed chunk", result.get(0).getContent());
|
||||
Assert.assertEquals(String.valueOf(knowledgeId), searcher.lastKnowledgeId);
|
||||
Assert.assertEquals(
|
||||
tech.easyflow.ai.entity.DocumentCollection.TYPE_DOCUMENT,
|
||||
result.get(0).getMetadata("resultType")
|
||||
);
|
||||
Assert.assertEquals(
|
||||
completedDocumentId,
|
||||
result.get(0).getMetadata("documentId")
|
||||
);
|
||||
Assert.assertEquals(
|
||||
completedDocument.getTitle(),
|
||||
result.get(0).getMetadata("sourceFileName")
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证 FAQ 检索使用当前数据库记录回填稳定的 FAQ 来源信息。
|
||||
*
|
||||
* @throws Exception 反射注入异常
|
||||
*/
|
||||
@Test
|
||||
public void searchShouldFillFaqSourceMetadataFromDatabase() throws Exception {
|
||||
BigInteger knowledgeId = BigInteger.ONE;
|
||||
BigInteger faqId = BigInteger.valueOf(2001);
|
||||
BigInteger categoryId = BigInteger.valueOf(3001);
|
||||
tech.easyflow.ai.entity.DocumentCollection collection =
|
||||
new tech.easyflow.ai.entity.DocumentCollection();
|
||||
collection.setId(knowledgeId);
|
||||
collection.setCollectionType(
|
||||
tech.easyflow.ai.entity.DocumentCollection.TYPE_FAQ
|
||||
);
|
||||
collection.setOptions(new HashMap<String, Object>() {{
|
||||
put(KEY_DOC_RECALL_MAX_NUM, 5);
|
||||
put(KEY_SIMILARITY_THRESHOLD, BigDecimal.ZERO);
|
||||
}});
|
||||
|
||||
tech.easyflow.ai.entity.FaqItem faqItem =
|
||||
new tech.easyflow.ai.entity.FaqItem();
|
||||
faqItem.setId(faqId);
|
||||
faqItem.setCollectionId(knowledgeId);
|
||||
faqItem.setCategoryId(categoryId);
|
||||
faqItem.setQuestion("如何申请账号？");
|
||||
faqItem.setAnswerText("请联系管理员。");
|
||||
TestKeywordSearcher searcher = new TestKeywordSearcher(
|
||||
List.of(buildHit(faqId, 0.85D))
|
||||
);
|
||||
|
||||
DocumentCollectionServiceImpl service =
|
||||
new TestDocumentCollectionService(collection);
|
||||
setField(
|
||||
service,
|
||||
"searcherFactory",
|
||||
new SearcherFactory(
|
||||
new StaticObjectProvider<DocumentSearcher>(searcher)
|
||||
)
|
||||
);
|
||||
setField(service, "faqItemMapper", mockFaqItemMapper(faqItem));
|
||||
|
||||
tech.easyflow.ai.rag.KnowledgeRetrievalRequest request =
|
||||
new tech.easyflow.ai.rag.KnowledgeRetrievalRequest();
|
||||
request.setKnowledgeId(knowledgeId);
|
||||
request.setQuery("账号");
|
||||
request.setRetrievalMode(
|
||||
com.easyagents.rag.retrieval.RetrievalMode.KEYWORD
|
||||
);
|
||||
|
||||
List<Document> result = service.search(request);
|
||||
|
||||
Assert.assertEquals(1, result.size());
|
||||
Document item = result.get(0);
|
||||
Assert.assertEquals(
|
||||
tech.easyflow.ai.entity.DocumentCollection.TYPE_FAQ,
|
||||
item.getMetadata("resultType")
|
||||
);
|
||||
Assert.assertEquals(faqId, item.getMetadata("faqId"));
|
||||
Assert.assertEquals(
|
||||
faqItem.getQuestion(),
|
||||
item.getMetadata("question")
|
||||
);
|
||||
Assert.assertEquals(
|
||||
faqItem.getAnswerText(),
|
||||
item.getMetadata("answerText")
|
||||
);
|
||||
Assert.assertEquals(categoryId, item.getMetadata("categoryId"));
|
||||
}
|
||||
|
||||
private static Document buildHit(BigInteger id, double score) {
|
||||
@@ -132,6 +216,27 @@ public class DocumentCollectionServiceImplTest {
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 创建返回固定 FAQ 的 Mapper 桩。
|
||||
*
|
||||
* @param faqItem FAQ 数据
|
||||
* @return Mapper 桩
|
||||
*/
|
||||
private static FaqItemMapper mockFaqItemMapper(
|
||||
tech.easyflow.ai.entity.FaqItem faqItem
|
||||
) {
|
||||
return (FaqItemMapper) Proxy.newProxyInstance(
|
||||
FaqItemMapper.class.getClassLoader(),
|
||||
new Class<?>[]{FaqItemMapper.class},
|
||||
(proxy, method, args) -> {
|
||||
if ("selectListByQuery".equals(method.getName())) {
|
||||
return List.of(faqItem);
|
||||
}
|
||||
return defaultValue(method.getReturnType());
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
private static void setField(Object target, String fieldName, Object value) throws Exception {
|
||||
Field field = DocumentCollectionServiceImpl.class.getDeclaredField(fieldName);
|
||||
field.setAccessible(true);
|
||||
|
||||
@@ -0,0 +1,255 @@
|
||||
package tech.easyflow.ai.service.impl;
|
||||
|
||||
import com.easyagents.core.store.DocumentStore;
|
||||
import com.easyagents.core.store.StoreResult;
|
||||
import com.mybatisflex.core.query.QueryWrapper;
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
import org.mockito.Mockito;
|
||||
import tech.easyflow.ai.config.SearcherFactory;
|
||||
import tech.easyflow.ai.entity.Document;
|
||||
import tech.easyflow.ai.entity.DocumentCollection;
|
||||
import tech.easyflow.ai.entity.Model;
|
||||
import tech.easyflow.ai.enums.DocumentProcessStatus;
|
||||
import tech.easyflow.ai.mapper.DocumentChunkMapper;
|
||||
import tech.easyflow.ai.mapper.DocumentMapper;
|
||||
import tech.easyflow.ai.service.DocumentCollectionService;
|
||||
import tech.easyflow.ai.service.ModelService;
|
||||
import tech.easyflow.common.filestorage.FileStorageService;
|
||||
import tech.easyflow.common.web.exceptions.BusinessException;
|
||||
|
||||
import java.lang.reflect.Field;
|
||||
import java.math.BigInteger;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* {@link DocumentServiceImpl} 文档维护回归测试。
|
||||
*
|
||||
* @author Codex
|
||||
* @since 2026-08-02
|
||||
*/
|
||||
public class DocumentServiceImplTest {
|
||||
|
||||
/**
|
||||
* 验证删除链路在外部索引和分块清理后删除文档主记录。
|
||||
*
|
||||
* @throws Exception 反射注入异常
|
||||
*/
|
||||
@Test
|
||||
public void removeDocShouldDeleteDocumentRecord() throws Exception {
|
||||
BigInteger documentId = BigInteger.valueOf(101);
|
||||
BigInteger knowledgeId = BigInteger.valueOf(102);
|
||||
BigInteger modelId = BigInteger.valueOf(103);
|
||||
Document document = new Document();
|
||||
document.setId(documentId);
|
||||
document.setCollectionId(knowledgeId);
|
||||
document.setDocumentPath("storage://document.txt");
|
||||
document.setProcessStatus(DocumentProcessStatus.COMPLETED.name());
|
||||
|
||||
DocumentStore documentStore = Mockito.mock(DocumentStore.class);
|
||||
Mockito.when(documentStore.delete(
|
||||
Mockito.anyList(), Mockito.any())).thenReturn(StoreResult.success());
|
||||
DocumentCollection knowledge = Mockito.mock(DocumentCollection.class);
|
||||
Mockito.when(knowledge.getVectorEmbedModelId()).thenReturn(modelId);
|
||||
Mockito.when(knowledge.getVectorStoreCollection()).thenReturn("kb-test");
|
||||
Mockito.when(knowledge.toDocumentStore()).thenReturn(documentStore);
|
||||
Model model = new Model();
|
||||
model.setModelName("embedding-test");
|
||||
|
||||
DocumentMapper documentMapper = Mockito.mock(DocumentMapper.class);
|
||||
Mockito.when(documentMapper.selectOneByQuery(
|
||||
Mockito.any(QueryWrapper.class))).thenReturn(document);
|
||||
Mockito.when(documentMapper.deleteById(documentId)).thenReturn(1);
|
||||
DocumentChunkMapper chunkMapper =
|
||||
Mockito.mock(DocumentChunkMapper.class);
|
||||
Mockito.when(chunkMapper.selectListByQueryAs(
|
||||
Mockito.any(QueryWrapper.class), Mockito.eq(BigInteger.class)))
|
||||
.thenReturn(List.of(BigInteger.valueOf(201)));
|
||||
Mockito.when(chunkMapper.deleteByQuery(
|
||||
Mockito.any(QueryWrapper.class))).thenReturn(1);
|
||||
DocumentCollectionService knowledgeService =
|
||||
Mockito.mock(DocumentCollectionService.class);
|
||||
Mockito.when(knowledgeService.getById(knowledgeId))
|
||||
.thenReturn(knowledge);
|
||||
ModelService modelService = Mockito.mock(ModelService.class);
|
||||
Mockito.when(modelService.getById(modelId)).thenReturn(model);
|
||||
FileStorageService storageService =
|
||||
Mockito.mock(FileStorageService.class);
|
||||
SearcherFactory searcherFactory = Mockito.mock(SearcherFactory.class);
|
||||
|
||||
DocumentServiceImpl service = new DocumentServiceImpl();
|
||||
setField(service, "documentMapper", documentMapper);
|
||||
setField(service, "documentChunkMapper", chunkMapper);
|
||||
setField(service, "knowledgeService", knowledgeService);
|
||||
setField(service, "modelService", modelService);
|
||||
setField(service, "storageService", storageService);
|
||||
setField(service, "searcherFactory", searcherFactory);
|
||||
|
||||
Assert.assertTrue(service.removeDoc(documentId.toString()));
|
||||
|
||||
Mockito.verify(documentMapper).deleteById(documentId);
|
||||
Mockito.verify(storageService).delete(document.getDocumentPath());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证最后一个分块已单独删除时，文档删除跳过空向量请求并完成清理。
|
||||
*
|
||||
* @throws Exception 反射注入异常
|
||||
*/
|
||||
@Test
|
||||
public void removeDocShouldHandleDocumentWithoutChunks() throws Exception {
|
||||
BigInteger documentId = BigInteger.valueOf(202);
|
||||
BigInteger knowledgeId = BigInteger.valueOf(203);
|
||||
BigInteger modelId = BigInteger.valueOf(204);
|
||||
Document document = new Document();
|
||||
document.setId(documentId);
|
||||
document.setCollectionId(knowledgeId);
|
||||
document.setDocumentPath("storage://empty-document.txt");
|
||||
document.setProcessStatus(DocumentProcessStatus.COMPLETED.name());
|
||||
|
||||
DocumentCollection knowledge = Mockito.mock(DocumentCollection.class);
|
||||
Mockito.when(knowledge.getVectorEmbedModelId()).thenReturn(modelId);
|
||||
Mockito.when(knowledge.getVectorStoreCollection()).thenReturn("kb-test");
|
||||
Mockito.when(knowledge.toDocumentStore()).thenReturn(null);
|
||||
DocumentMapper documentMapper = Mockito.mock(DocumentMapper.class);
|
||||
Mockito.when(documentMapper.selectOneByQuery(
|
||||
Mockito.any(QueryWrapper.class))).thenReturn(document);
|
||||
Mockito.when(documentMapper.deleteById(documentId)).thenReturn(1);
|
||||
DocumentChunkMapper chunkMapper = Mockito.mock(DocumentChunkMapper.class);
|
||||
Mockito.when(chunkMapper.selectListByQueryAs(
|
||||
Mockito.any(QueryWrapper.class), Mockito.eq(BigInteger.class)))
|
||||
.thenReturn(List.of());
|
||||
Mockito.when(chunkMapper.deleteByQuery(
|
||||
Mockito.any(QueryWrapper.class))).thenReturn(0);
|
||||
DocumentCollectionService knowledgeService =
|
||||
Mockito.mock(DocumentCollectionService.class);
|
||||
Mockito.when(knowledgeService.getById(knowledgeId)).thenReturn(knowledge);
|
||||
ModelService modelService = Mockito.mock(ModelService.class);
|
||||
FileStorageService storageService = Mockito.mock(FileStorageService.class);
|
||||
SearcherFactory searcherFactory = Mockito.mock(SearcherFactory.class);
|
||||
|
||||
DocumentServiceImpl service = new DocumentServiceImpl();
|
||||
setField(service, "documentMapper", documentMapper);
|
||||
setField(service, "documentChunkMapper", chunkMapper);
|
||||
setField(service, "knowledgeService", knowledgeService);
|
||||
setField(service, "modelService", modelService);
|
||||
setField(service, "storageService", storageService);
|
||||
setField(service, "searcherFactory", searcherFactory);
|
||||
|
||||
Assert.assertTrue(service.removeDoc(documentId.toString()));
|
||||
|
||||
Mockito.verify(knowledge, Mockito.never()).toDocumentStore();
|
||||
Mockito.verifyNoInteractions(modelService);
|
||||
Mockito.verify(documentMapper).deleteById(documentId);
|
||||
Mockito.verify(storageService).delete(document.getDocumentPath());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证分块中的文档禁止删除，且不会触发任何外部或数据库清理。
|
||||
*
|
||||
* @throws Exception 反射注入异常
|
||||
*/
|
||||
@Test
|
||||
public void removeDocShouldRejectSplittingDocumentWithoutSideEffects()
|
||||
throws Exception {
|
||||
BigInteger documentId = BigInteger.valueOf(301);
|
||||
Document document = new Document();
|
||||
document.setId(documentId);
|
||||
document.setProcessStatus(DocumentProcessStatus.SPLITTING.name());
|
||||
DocumentMapper documentMapper = Mockito.mock(DocumentMapper.class);
|
||||
Mockito.when(documentMapper.selectOneByQuery(
|
||||
Mockito.any(QueryWrapper.class))).thenReturn(document);
|
||||
DocumentChunkMapper chunkMapper = Mockito.mock(DocumentChunkMapper.class);
|
||||
DocumentCollectionService knowledgeService =
|
||||
Mockito.mock(DocumentCollectionService.class);
|
||||
|
||||
DocumentServiceImpl service = new DocumentServiceImpl();
|
||||
setField(service, "documentMapper", documentMapper);
|
||||
setField(service, "documentChunkMapper", chunkMapper);
|
||||
setField(service, "knowledgeService", knowledgeService);
|
||||
|
||||
try {
|
||||
service.removeDoc(documentId.toString());
|
||||
Assert.fail("分块中的文档应拒绝删除");
|
||||
} catch (BusinessException expected) {
|
||||
Assert.assertEquals("文档处理中，暂不允许删除", expected.getMessage());
|
||||
}
|
||||
|
||||
Mockito.verifyNoInteractions(chunkMapper, knowledgeService);
|
||||
Mockito.verify(documentMapper, Mockito.never()).deleteById(Mockito.any());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证向量删除失败时停止后续清理，避免接口返回虚假成功。
|
||||
*
|
||||
* @throws Exception 反射注入异常
|
||||
*/
|
||||
@Test
|
||||
public void removeDocShouldStopWhenVectorDeleteFails() throws Exception {
|
||||
BigInteger documentId = BigInteger.valueOf(401);
|
||||
BigInteger knowledgeId = BigInteger.valueOf(402);
|
||||
BigInteger modelId = BigInteger.valueOf(403);
|
||||
Document document = new Document();
|
||||
document.setId(documentId);
|
||||
document.setCollectionId(knowledgeId);
|
||||
document.setProcessStatus(DocumentProcessStatus.COMPLETED.name());
|
||||
DocumentStore documentStore = Mockito.mock(DocumentStore.class);
|
||||
Mockito.when(documentStore.delete(
|
||||
Mockito.anyList(), Mockito.any())).thenReturn(StoreResult.fail("测试失败"));
|
||||
DocumentCollection knowledge = Mockito.mock(DocumentCollection.class);
|
||||
Mockito.when(knowledge.getVectorEmbedModelId()).thenReturn(modelId);
|
||||
Mockito.when(knowledge.getVectorStoreCollection()).thenReturn("kb-test");
|
||||
Mockito.when(knowledge.toDocumentStore()).thenReturn(documentStore);
|
||||
Model model = new Model();
|
||||
model.setModelName("embedding-test");
|
||||
DocumentMapper documentMapper = Mockito.mock(DocumentMapper.class);
|
||||
Mockito.when(documentMapper.selectOneByQuery(
|
||||
Mockito.any(QueryWrapper.class))).thenReturn(document);
|
||||
DocumentChunkMapper chunkMapper = Mockito.mock(DocumentChunkMapper.class);
|
||||
Mockito.when(chunkMapper.selectListByQueryAs(
|
||||
Mockito.any(QueryWrapper.class), Mockito.eq(BigInteger.class)))
|
||||
.thenReturn(List.of(BigInteger.valueOf(404)));
|
||||
DocumentCollectionService knowledgeService =
|
||||
Mockito.mock(DocumentCollectionService.class);
|
||||
Mockito.when(knowledgeService.getById(knowledgeId)).thenReturn(knowledge);
|
||||
ModelService modelService = Mockito.mock(ModelService.class);
|
||||
Mockito.when(modelService.getById(modelId)).thenReturn(model);
|
||||
FileStorageService storageService = Mockito.mock(FileStorageService.class);
|
||||
SearcherFactory searcherFactory = Mockito.mock(SearcherFactory.class);
|
||||
|
||||
DocumentServiceImpl service = new DocumentServiceImpl();
|
||||
setField(service, "documentMapper", documentMapper);
|
||||
setField(service, "documentChunkMapper", chunkMapper);
|
||||
setField(service, "knowledgeService", knowledgeService);
|
||||
setField(service, "modelService", modelService);
|
||||
setField(service, "storageService", storageService);
|
||||
setField(service, "searcherFactory", searcherFactory);
|
||||
|
||||
try {
|
||||
service.removeDoc(documentId.toString());
|
||||
Assert.fail("向量删除失败时应停止文档删除");
|
||||
} catch (BusinessException expected) {
|
||||
Assert.assertEquals("文档向量删除失败", expected.getMessage());
|
||||
}
|
||||
|
||||
Mockito.verify(chunkMapper, Mockito.never()).deleteByQuery(Mockito.any());
|
||||
Mockito.verify(documentMapper, Mockito.never()).deleteById(Mockito.any());
|
||||
Mockito.verifyNoInteractions(storageService);
|
||||
}
|
||||
|
||||
/**
|
||||
* 通过反射注入测试依赖。
|
||||
*
|
||||
* @param target 目标对象
|
||||
* @param fieldName 字段名
|
||||
* @param value 字段值
|
||||
* @throws Exception 字段不存在或无法访问时抛出
|
||||
*/
|
||||
private static void setField(Object target,
|
||||
String fieldName,
|
||||
Object value) throws Exception {
|
||||
Field field = DocumentServiceImpl.class.getDeclaredField(fieldName);
|
||||
field.setAccessible(true);
|
||||
field.set(target, value);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package tech.easyflow.ai.service.impl;
|
||||
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
import org.mockito.ArgumentCaptor;
|
||||
import org.mockito.Mockito;
|
||||
import tech.easyflow.ai.enums.KnowledgeApiPermissionScope;
|
||||
import tech.easyflow.common.cache.RedisLockExecutor;
|
||||
import tech.easyflow.system.entity.SysApiKey;
|
||||
import tech.easyflow.system.entity.SysApiKeyResource;
|
||||
import tech.easyflow.system.entity.SysApiKeyResourceMapping;
|
||||
import tech.easyflow.system.service.SysApiKeyResourceMappingService;
|
||||
import tech.easyflow.system.service.SysApiKeyResourceService;
|
||||
import tech.easyflow.system.service.SysApiKeyService;
|
||||
|
||||
import java.lang.reflect.Field;
|
||||
import java.math.BigInteger;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.atomic.AtomicLong;
|
||||
|
||||
/**
|
||||
* 知识库 Public API 产品权限映射测试。
|
||||
*
|
||||
* @author Codex
|
||||
* @since 2026-08-02
|
||||
*/
|
||||
public class KnowledgeSharePermissionServiceImplTest {
|
||||
|
||||
/**
|
||||
* 验证读取、导入和维护权限分别只生成各自接口映射。
|
||||
*
|
||||
* @throws Exception 依赖注入失败
|
||||
*/
|
||||
@Test
|
||||
public void replacePermissionsShouldKeepReadImportMaintenanceSeparated()
|
||||
throws Exception {
|
||||
BigInteger apiKeyId = BigInteger.valueOf(701);
|
||||
SysApiKeyService apiKeyService = Mockito.mock(SysApiKeyService.class);
|
||||
SysApiKeyResourceService resourceService =
|
||||
Mockito.mock(SysApiKeyResourceService.class);
|
||||
SysApiKeyResourceMappingService mappingService =
|
||||
Mockito.mock(SysApiKeyResourceMappingService.class);
|
||||
RedisLockExecutor redisLockExecutor =
|
||||
Mockito.mock(RedisLockExecutor.class);
|
||||
RedisLockExecutor.LockHandle lockHandle =
|
||||
Mockito.mock(RedisLockExecutor.LockHandle.class);
|
||||
Mockito.when(redisLockExecutor.tryAcquire(
|
||||
Mockito.anyString(),
|
||||
Mockito.any(),
|
||||
Mockito.any()
|
||||
)).thenReturn(lockHandle);
|
||||
Mockito.when(apiKeyService.getById(apiKeyId)).thenReturn(new SysApiKey());
|
||||
AtomicLong resourceId = new AtomicLong(800);
|
||||
Mockito.when(resourceService.getOne(
|
||||
Mockito.any(com.mybatisflex.core.query.QueryWrapper.class)))
|
||||
.thenAnswer(invocation -> {
|
||||
SysApiKeyResource resource = new SysApiKeyResource();
|
||||
resource.setId(BigInteger.valueOf(resourceId.incrementAndGet()));
|
||||
return resource;
|
||||
});
|
||||
|
||||
KnowledgeSharePermissionServiceImpl service =
|
||||
new KnowledgeSharePermissionServiceImpl();
|
||||
setField(service, "sysApiKeyService", apiKeyService);
|
||||
setField(service, "resourceService", resourceService);
|
||||
setField(service, "mappingService", mappingService);
|
||||
setField(service, "redisLockExecutor", redisLockExecutor);
|
||||
|
||||
service.replaceApiPermissions(apiKeyId, true, false, false);
|
||||
service.replaceApiPermissions(apiKeyId, false, true, false);
|
||||
service.replaceApiPermissions(apiKeyId, false, false, true);
|
||||
|
||||
@SuppressWarnings("unchecked")
|
||||
ArgumentCaptor<List<SysApiKeyResourceMapping>> mappings =
|
||||
ArgumentCaptor.forClass(List.class);
|
||||
Mockito.verify(mappingService, Mockito.times(3)).remove(
|
||||
Mockito.any(com.mybatisflex.core.query.QueryWrapper.class)
|
||||
);
|
||||
Mockito.verify(mappingService, Mockito.times(3))
|
||||
.saveBatch(mappings.capture());
|
||||
Mockito.verify(lockHandle, Mockito.times(3)).release();
|
||||
Assert.assertEquals(3, mappings.getAllValues().size());
|
||||
assertScopeMappings(
|
||||
mappings.getAllValues().get(0),
|
||||
KnowledgeApiPermissionScope.KNOWLEDGE_READ,
|
||||
8
|
||||
);
|
||||
assertScopeMappings(
|
||||
mappings.getAllValues().get(1),
|
||||
KnowledgeApiPermissionScope.KNOWLEDGE_IMPORT,
|
||||
14
|
||||
);
|
||||
assertScopeMappings(
|
||||
mappings.getAllValues().get(2),
|
||||
KnowledgeApiPermissionScope.KNOWLEDGE_MAINTENANCE,
|
||||
6
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 断言一组全局知识库接口映射只包含指定权限。
|
||||
*
|
||||
* @param mappings 接口映射
|
||||
* @param expectedScope 预期权限
|
||||
* @param expectedSize 预期接口数量
|
||||
*/
|
||||
private void assertScopeMappings(
|
||||
List<SysApiKeyResourceMapping> mappings,
|
||||
KnowledgeApiPermissionScope expectedScope,
|
||||
int expectedSize) {
|
||||
Assert.assertEquals(expectedSize, mappings.size());
|
||||
for (SysApiKeyResourceMapping mapping : mappings) {
|
||||
Assert.assertEquals(
|
||||
expectedScope.name(),
|
||||
mapping.getActionScope()
|
||||
);
|
||||
Assert.assertNull(mapping.getResourceTargetId());
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 反射注入测试依赖。
|
||||
*
|
||||
* @param target 目标对象
|
||||
* @param fieldName 字段名
|
||||
* @param value 字段值
|
||||
* @throws Exception 字段不存在或不可访问
|
||||
*/
|
||||
private static void setField(
|
||||
Object target,
|
||||
String fieldName,
|
||||
Object value) throws Exception {
|
||||
Field field = target.getClass().getDeclaredField(fieldName);
|
||||
field.setAccessible(true);
|
||||
field.set(target, value);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user