feat: 支持知识库 CSV 大文件导入
- 增加 CSV 流式解析、表格语义分块和分页预览 - 增加快照两阶段清理、失败重试和格式校验 - 补充批量入口、管理端交互和回归测试
This commit is contained in:
@@ -0,0 +1,67 @@
|
||||
package tech.easyflow.ai.documentimport.task;
|
||||
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
|
||||
import java.lang.reflect.Field;
|
||||
import java.lang.reflect.Method;
|
||||
import java.util.Set;
|
||||
|
||||
/**
|
||||
* CSV 导入入口支持策略测试。
|
||||
*/
|
||||
public class CsvImportSupportPolicyTest {
|
||||
|
||||
/**
|
||||
* 验证单文件、管理端批次和 Public API 批次三条入口都允许 CSV。
|
||||
*
|
||||
* @throws Exception 反射读取失败
|
||||
*/
|
||||
@Test
|
||||
public void shouldAllowCsvAcrossAllKnowledgeImportEntrypoints()
|
||||
throws Exception {
|
||||
KnowledgeDocumentImportTaskAppService taskService =
|
||||
new KnowledgeDocumentImportTaskAppService();
|
||||
Method assertSupported = KnowledgeDocumentImportTaskAppService.class
|
||||
.getDeclaredMethod("assertSupportedImportFile", String.class);
|
||||
assertSupported.setAccessible(true);
|
||||
Method normalizeExtension =
|
||||
KnowledgeDocumentImportTaskAppService.class.getDeclaredMethod(
|
||||
"normalizeFileExtension", String.class, String.class);
|
||||
normalizeExtension.setAccessible(true);
|
||||
String lowerCaseExtension = (String) normalizeExtension.invoke(
|
||||
taskService, "results.csv", "/knowledge/results.csv");
|
||||
String upperCaseExtension = (String) normalizeExtension.invoke(
|
||||
taskService, "RESULTS.CSV", "/knowledge/RESULTS.CSV");
|
||||
Assert.assertEquals("csv", lowerCaseExtension);
|
||||
Assert.assertEquals("csv", upperCaseExtension);
|
||||
assertSupported.invoke(taskService, lowerCaseExtension);
|
||||
assertSupported.invoke(taskService, upperCaseExtension);
|
||||
|
||||
Assert.assertTrue(readSupportedExtensions(
|
||||
DocumentImportBatchAppService.class).contains("csv"));
|
||||
Assert.assertTrue(readSupportedExtensions(
|
||||
KnowledgeImportBatchFacade.class).contains("csv"));
|
||||
Assert.assertSame(
|
||||
DocumentImportFormatPolicy.supportedExtensions(),
|
||||
readSupportedExtensions(DocumentImportBatchAppService.class));
|
||||
Assert.assertSame(
|
||||
DocumentImportFormatPolicy.supportedExtensions(),
|
||||
readSupportedExtensions(KnowledgeImportBatchFacade.class));
|
||||
}
|
||||
|
||||
/**
|
||||
* 读取入口类的支持扩展名集合。
|
||||
*
|
||||
* @param type 入口类
|
||||
* @return 扩展名集合
|
||||
* @throws Exception 字段读取失败
|
||||
*/
|
||||
@SuppressWarnings("unchecked")
|
||||
private Set<String> readSupportedExtensions(Class<?> type)
|
||||
throws Exception {
|
||||
Field field = type.getDeclaredField("SUPPORTED_EXTENSIONS");
|
||||
field.setAccessible(true);
|
||||
return (Set<String>) field.get(null);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,361 @@
|
||||
package tech.easyflow.ai.documentimport.task;
|
||||
|
||||
import com.easyagents.rag.core.BgeM3ChunkSafety;
|
||||
import com.easyagents.rag.ingestion.model.StrategyConfig;
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
import tech.easyflow.ai.documentimport.DocumentImportDtos;
|
||||
import tech.easyflow.ai.entity.DocumentChunk;
|
||||
|
||||
import java.lang.reflect.Field;
|
||||
import java.math.BigInteger;
|
||||
import java.nio.charset.Charset;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* {@link CsvTableSnapshotService} 流式解析与分块测试。
|
||||
*/
|
||||
public class CsvTableSnapshotServiceTest {
|
||||
|
||||
/**
|
||||
* 验证引号逗号、跨行字段、尾空列和重复表头都能稳定解析并按行窗口分块。
|
||||
*
|
||||
* @throws Exception 反射注入失败
|
||||
*/
|
||||
@Test
|
||||
public void shouldParseRfc4180AndBuildRowWindowChunks() throws Exception {
|
||||
InMemoryFileStorageService storage = new InMemoryFileStorageService();
|
||||
String sourcePath = "source/demo.csv";
|
||||
String csv = "姓名,备注,姓名,\r\n"
|
||||
+ "张三,\"第一行\n第二行\",A,\r\n"
|
||||
+ "李四,\"含,逗号\",B,\r\n";
|
||||
storage.put(sourcePath, csv.getBytes(StandardCharsets.UTF_8));
|
||||
CsvTableSnapshotService service = createService(storage);
|
||||
|
||||
CsvTableSnapshotService.CsvParseResult result = service.parse(
|
||||
sourcePath, BigInteger.valueOf(7), BigInteger.valueOf(9), "token-1");
|
||||
|
||||
Assert.assertEquals("UTF-8", result.getEncoding());
|
||||
Assert.assertEquals(2L, result.getRowCount());
|
||||
Assert.assertEquals(4, result.getColumnCount());
|
||||
Assert.assertEquals(
|
||||
List.of("姓名", "备注", "姓名_2", "列_4"), result.getHeaders());
|
||||
|
||||
tech.easyflow.ai.entity.Document document =
|
||||
new tech.easyflow.ai.entity.Document();
|
||||
document.setId(BigInteger.valueOf(9));
|
||||
document.setCollectionId(BigInteger.valueOf(7));
|
||||
document.setTitle("demo.csv");
|
||||
document.setDocumentPath(sourcePath);
|
||||
StrategyConfig strategy = StrategyConfig.defaults();
|
||||
strategy.setStrategyCode("TABLE_ROW");
|
||||
strategy.setRowsPerChunk(1);
|
||||
|
||||
DocumentImportDtos.PreviewSession session =
|
||||
service.buildChunkSnapshot(
|
||||
document, result.getManifestPath(), strategy);
|
||||
|
||||
Assert.assertEquals(Integer.valueOf(2), session.getTotalChunks());
|
||||
Assert.assertEquals(2, session.getDocumentChunks().size());
|
||||
Assert.assertTrue(session.getDocumentChunks().get(0).getContent()
|
||||
.contains("第一行<br/>第二行"));
|
||||
Assert.assertTrue(session.getDocumentChunks().get(1).getContent()
|
||||
.contains("含,逗号"));
|
||||
Assert.assertTrue(session.getDocumentChunks().stream()
|
||||
.allMatch(chunk -> "TABLE_ROW".equals(
|
||||
chunk.getOptions().get("chunkType"))));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证无 BOM 且非 UTF-8 的 CSV 会重新打开源文件并回退到 GB18030。
|
||||
*
|
||||
* @throws Exception 反射注入失败
|
||||
*/
|
||||
@Test
|
||||
public void shouldFallbackToGb18030AfterStrictUtf8Failure() throws Exception {
|
||||
InMemoryFileStorageService storage = new InMemoryFileStorageService();
|
||||
String sourcePath = "source/gb.csv";
|
||||
storage.put(
|
||||
sourcePath,
|
||||
"名称,说明\r\n测试,中文内容\r\n".getBytes(Charset.forName("GB18030")));
|
||||
CsvTableSnapshotService service = createService(storage);
|
||||
|
||||
CsvTableSnapshotService.CsvParseResult result = service.parse(
|
||||
sourcePath, BigInteger.valueOf(7), BigInteger.valueOf(10), "token-2");
|
||||
|
||||
Assert.assertEquals("GB18030", result.getEncoding());
|
||||
Assert.assertEquals(1L, result.getRowCount());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证数据行列数与表头不一致时返回稳定失败码。
|
||||
*
|
||||
* @throws Exception 反射注入失败
|
||||
*/
|
||||
@Test
|
||||
public void shouldRejectColumnMismatchWithStableFailureCode() throws Exception {
|
||||
InMemoryFileStorageService storage = new InMemoryFileStorageService();
|
||||
String sourcePath = "source/broken.csv";
|
||||
storage.put(
|
||||
sourcePath,
|
||||
"a,b\r\n1,2,3\r\n".getBytes(StandardCharsets.UTF_8));
|
||||
CsvTableSnapshotService service = createService(storage);
|
||||
|
||||
try {
|
||||
service.parse(
|
||||
sourcePath, BigInteger.valueOf(7), BigInteger.valueOf(11), "token-3");
|
||||
Assert.fail("应拒绝列数不一致的 CSV");
|
||||
} catch (CsvImportException error) {
|
||||
Assert.assertEquals(
|
||||
CsvTableSnapshotService.FAILURE_COLUMN_MISMATCH,
|
||||
error.getFailureCode());
|
||||
Assert.assertTrue(error.getMessage().contains("列数"));
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证二进制伪装 CSV 中的 NUL 字节会被明确拒绝。
|
||||
*
|
||||
* @throws Exception 反射注入失败
|
||||
*/
|
||||
@Test
|
||||
public void shouldRejectNulByteAsMalformedCsv() throws Exception {
|
||||
InMemoryFileStorageService storage = new InMemoryFileStorageService();
|
||||
String sourcePath = "source/binary.csv";
|
||||
storage.put(
|
||||
sourcePath,
|
||||
"name,value\nalice,\0binary\n".getBytes(StandardCharsets.UTF_8));
|
||||
CsvTableSnapshotService service = createService(storage);
|
||||
|
||||
try {
|
||||
service.parse(
|
||||
sourcePath, BigInteger.valueOf(7), BigInteger.valueOf(12), "token-4");
|
||||
Assert.fail("包含 NUL 字节的文件不应进入后续阶段");
|
||||
} catch (CsvImportException error) {
|
||||
Assert.assertEquals(
|
||||
CsvTableSnapshotService.FAILURE_MALFORMED,
|
||||
error.getFailureCode());
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证超长表格正文会生成多个不超过数据库上限的续片。
|
||||
*/
|
||||
@Test
|
||||
public void shouldSplitOversizedWindowWithoutSilentTruncation() {
|
||||
TabularRowWindowChunkBuilder builder =
|
||||
new TabularRowWindowChunkBuilder();
|
||||
String longValue = "长".repeat(
|
||||
TabularRowWindowChunkBuilder.MAX_CHUNK_CONTENT_CHARS + 100);
|
||||
|
||||
List<DocumentChunk> chunks = builder.build(
|
||||
BigInteger.ONE,
|
||||
BigInteger.valueOf(2),
|
||||
"大表",
|
||||
List.of("正文"),
|
||||
List.of(new TabularRowWindowChunkBuilder.TabularRow(
|
||||
2, List.of(longValue))),
|
||||
1,
|
||||
"TABLE_ROW");
|
||||
|
||||
Assert.assertTrue(chunks.size() > 1);
|
||||
Assert.assertTrue(chunks.stream().allMatch(
|
||||
chunk -> chunk.getContent().length()
|
||||
<= TabularRowWindowChunkBuilder.MAX_CHUNK_CONTENT_CHARS));
|
||||
Assert.assertTrue(chunks.stream().allMatch(
|
||||
chunk -> BgeM3ChunkSafety.isWithinHardLimit(chunk.getContent())));
|
||||
Assert.assertTrue(chunks.stream().allMatch(
|
||||
chunk -> chunk.getContent().startsWith("# 大表")
|
||||
&& chunk.getContent().contains("| 正文 |")));
|
||||
Assert.assertTrue(chunks.stream().allMatch(
|
||||
chunk -> Integer.valueOf(2).equals(
|
||||
chunk.getOptions().get("rowStart"))
|
||||
&& Integer.valueOf(2).equals(
|
||||
chunk.getOptions().get("rowEnd"))));
|
||||
long restoredChars = chunks.stream()
|
||||
.map(DocumentChunk::getContent)
|
||||
.flatMapToInt(String::chars)
|
||||
.filter(value -> value == '长')
|
||||
.count();
|
||||
Assert.assertEquals(longValue.length(), restoredChars);
|
||||
Assert.assertEquals(
|
||||
chunks.size(), chunks.get(0).getOptions().get("partTotal"));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证超长记录优先按列边界续片，且每个续片都保留标题和列名。
|
||||
*/
|
||||
@Test
|
||||
public void shouldSplitOversizedRowAtColumnBoundaries() {
|
||||
TabularRowWindowChunkBuilder builder =
|
||||
new TabularRowWindowChunkBuilder();
|
||||
String longValue = "长".repeat(20_000);
|
||||
|
||||
List<DocumentChunk> chunks = builder.build(
|
||||
BigInteger.ONE,
|
||||
BigInteger.valueOf(2),
|
||||
"列边界表",
|
||||
List.of("左列", "正文", "右列"),
|
||||
List.of(new TabularRowWindowChunkBuilder.TabularRow(
|
||||
8, List.of("LEFT_MARKER", longValue, "RIGHT_MARKER"))),
|
||||
1,
|
||||
"TABLE_ROW");
|
||||
|
||||
Assert.assertTrue(chunks.size() > 2);
|
||||
Assert.assertTrue(chunks.stream().allMatch(
|
||||
chunk -> chunk.getContent().startsWith("# 列边界表")
|
||||
&& chunk.getContent().contains("\n| ")));
|
||||
String allContent = chunks.stream()
|
||||
.map(DocumentChunk::getContent)
|
||||
.reduce("", String::concat);
|
||||
Assert.assertEquals(1, occurrences(allContent, "LEFT_MARKER"));
|
||||
Assert.assertEquals(1, occurrences(allContent, "RIGHT_MARKER"));
|
||||
Assert.assertEquals(
|
||||
longValue.length(),
|
||||
allContent.chars().filter(value -> value == '长').count());
|
||||
Assert.assertTrue(chunks.stream().allMatch(
|
||||
chunk -> Integer.valueOf(8).equals(
|
||||
chunk.getOptions().get("rowStart"))
|
||||
&& Integer.valueOf(8).equals(
|
||||
chunk.getOptions().get("rowEnd"))));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证行数未满时仍会在加入下一行超过 Token 上限前闭合窗口。
|
||||
*
|
||||
* @throws Exception 反射注入失败
|
||||
*/
|
||||
@Test
|
||||
public void shouldCloseWindowBeforeNextRowExceedsTokenLimit()
|
||||
throws Exception {
|
||||
InMemoryFileStorageService storage =
|
||||
new InMemoryFileStorageService();
|
||||
String sourcePath = "source/token-window.csv";
|
||||
String value = "中".repeat(5_000);
|
||||
storage.put(
|
||||
sourcePath,
|
||||
("正文\n" + value + "\n" + value + "\n")
|
||||
.getBytes(StandardCharsets.UTF_8));
|
||||
CsvTableSnapshotService service = createService(storage);
|
||||
CsvTableSnapshotService.CsvParseResult result = service.parse(
|
||||
sourcePath,
|
||||
BigInteger.valueOf(7),
|
||||
BigInteger.valueOf(13),
|
||||
"token-window");
|
||||
tech.easyflow.ai.entity.Document document =
|
||||
new tech.easyflow.ai.entity.Document();
|
||||
document.setId(BigInteger.valueOf(13));
|
||||
document.setCollectionId(BigInteger.valueOf(7));
|
||||
document.setTitle("token-window.csv");
|
||||
StrategyConfig strategy = StrategyConfig.defaults();
|
||||
strategy.setStrategyCode("TABLE_ROW");
|
||||
strategy.setRowsPerChunk(10);
|
||||
|
||||
DocumentImportDtos.PreviewSession session =
|
||||
service.buildChunkSnapshot(
|
||||
document, result.getManifestPath(), strategy);
|
||||
|
||||
Assert.assertEquals(Integer.valueOf(2), session.getTotalChunks());
|
||||
Assert.assertEquals(
|
||||
List.of(2, 3),
|
||||
session.getDocumentChunks().stream()
|
||||
.map(chunk -> (Integer) chunk.getOptions().get("rowStart"))
|
||||
.toList());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证 CSV 行分片删除失败时清单保持可重试。
|
||||
*
|
||||
* @throws Exception 反射注入失败
|
||||
*/
|
||||
@Test
|
||||
public void shouldKeepCsvManifestWhenPartDeletionFails()
|
||||
throws Exception {
|
||||
InMemoryFileStorageService storage =
|
||||
new InMemoryFileStorageService();
|
||||
String sourcePath = "source/cleanup.csv";
|
||||
storage.put(
|
||||
sourcePath,
|
||||
"a,b\n1,2\n".getBytes(StandardCharsets.UTF_8));
|
||||
CsvTableSnapshotService service = createService(storage);
|
||||
CsvTableSnapshotService.CsvParseResult result = service.parse(
|
||||
sourcePath,
|
||||
BigInteger.valueOf(7),
|
||||
BigInteger.valueOf(14),
|
||||
"cleanup");
|
||||
String manifestPath = result.getManifestPath();
|
||||
String partPath = storage.paths().stream()
|
||||
.filter(path -> !path.equals(sourcePath))
|
||||
.filter(path -> !path.equals(manifestPath))
|
||||
.findFirst()
|
||||
.orElseThrow();
|
||||
storage.failDelete(partPath, 1);
|
||||
|
||||
try {
|
||||
service.delete(manifestPath);
|
||||
Assert.fail("CSV 行分片删除失败时应抛出异常");
|
||||
} catch (IllegalStateException expected) {
|
||||
Assert.assertTrue(expected.getMessage().contains("模拟删除失败"));
|
||||
}
|
||||
Assert.assertTrue(storage.contains(manifestPath));
|
||||
Assert.assertTrue(storage.contains(partPath));
|
||||
|
||||
service.delete(manifestPath);
|
||||
Assert.assertEquals(1, storage.size());
|
||||
Assert.assertTrue(storage.contains(sourcePath));
|
||||
}
|
||||
|
||||
/**
|
||||
* 统计文本出现次数。
|
||||
*
|
||||
* @param content 完整文本
|
||||
* @param target 目标文本
|
||||
* @return 出现次数
|
||||
*/
|
||||
private int occurrences(String content, String target) {
|
||||
int count = 0;
|
||||
int offset = 0;
|
||||
while ((offset = content.indexOf(target, offset)) >= 0) {
|
||||
count++;
|
||||
offset += target.length();
|
||||
}
|
||||
return count;
|
||||
}
|
||||
|
||||
/**
|
||||
* 创建完成依赖注入的被测服务。
|
||||
*
|
||||
* @param storage 内存文件存储
|
||||
* @return CSV 服务
|
||||
* @throws Exception 反射注入失败
|
||||
*/
|
||||
private CsvTableSnapshotService createService(
|
||||
InMemoryFileStorageService storage) throws Exception {
|
||||
DocumentImportChunkSnapshotService chunkSnapshotService =
|
||||
new DocumentImportChunkSnapshotService();
|
||||
setField(chunkSnapshotService, "storageService", storage);
|
||||
|
||||
CsvTableSnapshotService service = new CsvTableSnapshotService();
|
||||
setField(service, "storageService", storage);
|
||||
setField(service, "chunkSnapshotService", chunkSnapshotService);
|
||||
return service;
|
||||
}
|
||||
|
||||
/**
|
||||
* 反射设置测试字段。
|
||||
*
|
||||
* @param target 目标对象
|
||||
* @param fieldName 字段名
|
||||
* @param value 字段值
|
||||
* @throws Exception 字段不存在或不可访问
|
||||
*/
|
||||
private void setField(
|
||||
Object target,
|
||||
String fieldName,
|
||||
Object value) throws Exception {
|
||||
Field field = target.getClass().getDeclaredField(fieldName);
|
||||
field.setAccessible(true);
|
||||
field.set(target, value);
|
||||
}
|
||||
}
|
||||
@@ -2,17 +2,13 @@ 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.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
|
||||
/**
|
||||
* {@link DocumentImportChunkSnapshotService} 持久化恢复测试。
|
||||
@@ -26,19 +22,8 @@ public class DocumentImportChunkSnapshotServiceTest {
|
||||
*/
|
||||
@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()));
|
||||
InMemoryFileStorageService storageService =
|
||||
new InMemoryFileStorageService();
|
||||
|
||||
DocumentImportChunkSnapshotService service = new DocumentImportChunkSnapshotService();
|
||||
Field storageField = DocumentImportChunkSnapshotService.class
|
||||
@@ -59,10 +44,108 @@ public class DocumentImportChunkSnapshotServiceTest {
|
||||
String path = service.save(session);
|
||||
DocumentImportDtos.PreviewSession restored = service.load(path);
|
||||
|
||||
Assert.assertEquals(storedPath, path);
|
||||
Assert.assertTrue(path.endsWith("-manifest.json"));
|
||||
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());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证 V2 快照可分页并按小批顺序消费，删除时同时清理清单和分片。
|
||||
*
|
||||
* @throws Exception 反射注入异常
|
||||
*/
|
||||
@Test
|
||||
public void shouldReadV2SnapshotByPageAndBatch() throws Exception {
|
||||
InMemoryFileStorageService storageService =
|
||||
new InMemoryFileStorageService();
|
||||
DocumentImportChunkSnapshotService service =
|
||||
new DocumentImportChunkSnapshotService();
|
||||
Field storageField = DocumentImportChunkSnapshotService.class
|
||||
.getDeclaredField("storageService");
|
||||
storageField.setAccessible(true);
|
||||
storageField.set(service, storageService);
|
||||
|
||||
DocumentImportDtos.PreviewSession session =
|
||||
new DocumentImportDtos.PreviewSession();
|
||||
session.setKnowledgeId(BigInteger.valueOf(7));
|
||||
session.setDocumentId(BigInteger.valueOf(9));
|
||||
List<DocumentChunk> chunks = new ArrayList<DocumentChunk>();
|
||||
for (int index = 0; index < 5; index++) {
|
||||
DocumentChunk chunk = new DocumentChunk();
|
||||
chunk.setId(BigInteger.valueOf(100 + index));
|
||||
chunk.setDocumentId(BigInteger.valueOf(9));
|
||||
chunk.setDocumentCollectionId(BigInteger.valueOf(7));
|
||||
chunk.setContent("chunk-" + index);
|
||||
chunks.add(chunk);
|
||||
}
|
||||
session.setDocumentChunks(chunks);
|
||||
|
||||
String path = service.save(session);
|
||||
List<DocumentChunk> page = service.loadPage(path, 2, 2);
|
||||
List<String> consumed = new ArrayList<String>();
|
||||
service.forEachBatch(path, 2, batch -> {
|
||||
Assert.assertTrue(batch.size() <= 2);
|
||||
for (DocumentChunk chunk : batch) {
|
||||
consumed.add(chunk.getContent());
|
||||
}
|
||||
});
|
||||
|
||||
Assert.assertEquals(List.of("chunk-2", "chunk-3"),
|
||||
page.stream().map(DocumentChunk::getContent).toList());
|
||||
Assert.assertEquals(
|
||||
List.of("chunk-0", "chunk-1", "chunk-2", "chunk-3", "chunk-4"),
|
||||
consumed);
|
||||
Assert.assertTrue(storageService.size() >= 2);
|
||||
service.delete(path);
|
||||
Assert.assertEquals(0, storageService.size());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证分片删除失败时保留清单，使后续重试仍能获取精确对象集合。
|
||||
*
|
||||
* @throws Exception 反射注入异常
|
||||
*/
|
||||
@Test
|
||||
public void shouldKeepManifestWhenPartDeletionFails() throws Exception {
|
||||
InMemoryFileStorageService storageService =
|
||||
new InMemoryFileStorageService();
|
||||
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(21));
|
||||
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 manifestPath = service.save(session);
|
||||
String partPath = storageService.paths().stream()
|
||||
.filter(path -> !path.equals(manifestPath))
|
||||
.findFirst()
|
||||
.orElseThrow();
|
||||
storageService.failDelete(partPath, 1);
|
||||
|
||||
try {
|
||||
service.delete(manifestPath);
|
||||
Assert.fail("分片删除失败时应抛出异常");
|
||||
} catch (IllegalStateException expected) {
|
||||
Assert.assertTrue(expected.getMessage().contains("模拟删除失败"));
|
||||
}
|
||||
Assert.assertTrue(storageService.contains(manifestPath));
|
||||
Assert.assertTrue(storageService.contains(partPath));
|
||||
|
||||
service.delete(manifestPath);
|
||||
Assert.assertEquals(0, storageService.size());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,237 @@
|
||||
package tech.easyflow.ai.documentimport.task;
|
||||
|
||||
import com.mybatisflex.core.query.QueryWrapper;
|
||||
import org.junit.Test;
|
||||
import org.mockito.Mockito;
|
||||
import tech.easyflow.ai.entity.DocumentImportSnapshotCleanup;
|
||||
import tech.easyflow.ai.mapper.DocumentImportSnapshotCleanupMapper;
|
||||
import tech.easyflow.ai.mapper.DocumentMapper;
|
||||
|
||||
import java.lang.reflect.Field;
|
||||
import java.math.BigInteger;
|
||||
import java.util.Date;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
|
||||
/**
|
||||
* {@link DocumentImportSnapshotCleanupService} 可靠清理测试。
|
||||
*/
|
||||
public class DocumentImportSnapshotCleanupServiceTest {
|
||||
|
||||
/**
|
||||
* 验证分片和清单按持久化阶段顺序删除。
|
||||
*
|
||||
* @throws Exception 反射注入失败
|
||||
*/
|
||||
@Test
|
||||
public void shouldPersistPhaseBeforeDeletingManifest()
|
||||
throws Exception {
|
||||
CleanupHarness harness = createHarness();
|
||||
|
||||
harness.service.scheduleChunkSnapshot(
|
||||
"snapshot/chunk-manifest.json");
|
||||
|
||||
Mockito.verify(harness.chunkSnapshotService).deleteParts(
|
||||
"snapshot/chunk-manifest.json");
|
||||
Mockito.verify(harness.cleanupMapper).advanceToManifest(
|
||||
Mockito.any(BigInteger.class),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class));
|
||||
Mockito.verify(harness.chunkSnapshotService).deleteManifest(
|
||||
"snapshot/chunk-manifest.json");
|
||||
Mockito.verify(harness.cleanupMapper).deleteCompleted(
|
||||
Mockito.any(BigInteger.class),
|
||||
Mockito.anyString());
|
||||
Mockito.verify(harness.cleanupMapper, Mockito.never())
|
||||
.releaseForRetry(
|
||||
Mockito.any(BigInteger.class),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证分片删除失败后保留记录并安排重试。
|
||||
*
|
||||
* @throws Exception 反射注入失败
|
||||
*/
|
||||
@Test
|
||||
public void shouldReleaseFailedCleanupForRetry()
|
||||
throws Exception {
|
||||
CleanupHarness harness = createHarness();
|
||||
Mockito.doThrow(new IllegalStateException("对象存储暂不可用"))
|
||||
.when(harness.chunkSnapshotService)
|
||||
.deleteParts("snapshot/failing-manifest.json");
|
||||
|
||||
harness.service.scheduleChunkSnapshot(
|
||||
"snapshot/failing-manifest.json");
|
||||
|
||||
Mockito.verify(harness.cleanupMapper).releaseForRetry(
|
||||
Mockito.any(BigInteger.class),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class),
|
||||
Mockito.contains("对象存储暂不可用"),
|
||||
Mockito.any(Date.class));
|
||||
Mockito.verify(harness.chunkSnapshotService, Mockito.never())
|
||||
.deleteManifest(Mockito.anyString());
|
||||
Mockito.verify(harness.cleanupMapper, Mockito.never())
|
||||
.deleteCompleted(
|
||||
Mockito.any(BigInteger.class),
|
||||
Mockito.anyString());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证 CSV 清单删除后原子清理文档中的旧快照指针。
|
||||
*
|
||||
* @throws Exception 反射注入失败
|
||||
*/
|
||||
@Test
|
||||
public void shouldClearCsvPointerAfterManifestDeletion()
|
||||
throws Exception {
|
||||
CleanupHarness harness = createHarness();
|
||||
BigInteger knowledgeId = BigInteger.valueOf(7);
|
||||
BigInteger documentId = BigInteger.valueOf(9);
|
||||
|
||||
harness.service.scheduleCsvTableSnapshot(
|
||||
knowledgeId,
|
||||
documentId,
|
||||
"snapshot/csv-manifest.json");
|
||||
|
||||
Mockito.verify(harness.csvTableSnapshotService).deleteParts(
|
||||
"snapshot/csv-manifest.json");
|
||||
Mockito.verify(harness.csvTableSnapshotService).deleteManifest(
|
||||
"snapshot/csv-manifest.json");
|
||||
Mockito.verify(harness.documentMapper).clearCsvSnapshotPath(
|
||||
Mockito.eq(documentId),
|
||||
Mockito.eq("snapshot/csv-manifest.json"),
|
||||
Mockito.any(Date.class));
|
||||
}
|
||||
|
||||
/**
|
||||
* 创建带内存状态的清理服务测试夹具。
|
||||
*
|
||||
* @return 测试夹具
|
||||
* @throws Exception 反射注入失败
|
||||
*/
|
||||
@SuppressWarnings("unchecked")
|
||||
private CleanupHarness createHarness() throws Exception {
|
||||
DocumentImportSnapshotCleanupService service =
|
||||
new DocumentImportSnapshotCleanupService();
|
||||
DocumentImportSnapshotCleanupMapper cleanupMapper =
|
||||
Mockito.mock(DocumentImportSnapshotCleanupMapper.class);
|
||||
DocumentImportChunkSnapshotService chunkSnapshotService =
|
||||
Mockito.mock(DocumentImportChunkSnapshotService.class);
|
||||
CsvTableSnapshotService csvTableSnapshotService =
|
||||
Mockito.mock(CsvTableSnapshotService.class);
|
||||
DocumentMapper documentMapper =
|
||||
Mockito.mock(DocumentMapper.class);
|
||||
AtomicReference<DocumentImportSnapshotCleanup> recordRef =
|
||||
new AtomicReference<DocumentImportSnapshotCleanup>();
|
||||
|
||||
Mockito.when(cleanupMapper.insertIgnore(
|
||||
Mockito.any(DocumentImportSnapshotCleanup.class)))
|
||||
.thenAnswer(invocation -> {
|
||||
recordRef.set(invocation.getArgument(0));
|
||||
return 1;
|
||||
});
|
||||
Mockito.when(cleanupMapper.selectOneByQuery(
|
||||
Mockito.any(QueryWrapper.class)))
|
||||
.thenAnswer(invocation -> recordRef.get());
|
||||
Mockito.when(cleanupMapper.claim(
|
||||
Mockito.any(BigInteger.class),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class),
|
||||
Mockito.any(Date.class)))
|
||||
.thenAnswer(invocation -> {
|
||||
DocumentImportSnapshotCleanup record = recordRef.get();
|
||||
record.setExecutionToken(invocation.getArgument(1));
|
||||
record.setAttemptCount(
|
||||
(record.getAttemptCount() == null
|
||||
? 0
|
||||
: record.getAttemptCount()) + 1);
|
||||
return 1;
|
||||
});
|
||||
Mockito.when(cleanupMapper.advanceToManifest(
|
||||
Mockito.any(BigInteger.class),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class)))
|
||||
.thenAnswer(invocation -> {
|
||||
recordRef.get().setPhase("MANIFEST_PENDING");
|
||||
return 1;
|
||||
});
|
||||
Mockito.when(cleanupMapper.deleteCompleted(
|
||||
Mockito.any(BigInteger.class),
|
||||
Mockito.anyString())).thenReturn(1);
|
||||
Mockito.when(cleanupMapper.releaseForRetry(
|
||||
Mockito.any(BigInteger.class),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class),
|
||||
Mockito.anyString(),
|
||||
Mockito.any(Date.class))).thenReturn(1);
|
||||
|
||||
setField(service, "cleanupMapper", cleanupMapper);
|
||||
setField(
|
||||
service, "chunkSnapshotService", chunkSnapshotService);
|
||||
setField(
|
||||
service, "csvTableSnapshotService", csvTableSnapshotService);
|
||||
setField(service, "documentMapper", documentMapper);
|
||||
return new CleanupHarness(
|
||||
service,
|
||||
cleanupMapper,
|
||||
chunkSnapshotService,
|
||||
csvTableSnapshotService,
|
||||
documentMapper);
|
||||
}
|
||||
|
||||
/**
|
||||
* 反射设置测试字段。
|
||||
*
|
||||
* @param target 目标对象
|
||||
* @param fieldName 字段名
|
||||
* @param value 字段值
|
||||
* @throws Exception 字段不存在或不可访问
|
||||
*/
|
||||
private void setField(
|
||||
Object target,
|
||||
String fieldName,
|
||||
Object value) throws Exception {
|
||||
Field field = target.getClass().getDeclaredField(fieldName);
|
||||
field.setAccessible(true);
|
||||
field.set(target, value);
|
||||
}
|
||||
|
||||
/**
|
||||
* 快照清理测试依赖集合。
|
||||
*/
|
||||
private static final class CleanupHarness {
|
||||
|
||||
private final DocumentImportSnapshotCleanupService service;
|
||||
private final DocumentImportSnapshotCleanupMapper cleanupMapper;
|
||||
private final DocumentImportChunkSnapshotService
|
||||
chunkSnapshotService;
|
||||
private final CsvTableSnapshotService csvTableSnapshotService;
|
||||
private final DocumentMapper documentMapper;
|
||||
|
||||
/**
|
||||
* 创建测试依赖集合。
|
||||
*
|
||||
* @param service 清理服务
|
||||
* @param cleanupMapper 清理 Mapper
|
||||
* @param chunkSnapshotService 分块快照服务
|
||||
* @param csvTableSnapshotService CSV 快照服务
|
||||
* @param documentMapper 文档 Mapper
|
||||
*/
|
||||
private CleanupHarness(
|
||||
DocumentImportSnapshotCleanupService service,
|
||||
DocumentImportSnapshotCleanupMapper cleanupMapper,
|
||||
DocumentImportChunkSnapshotService chunkSnapshotService,
|
||||
CsvTableSnapshotService csvTableSnapshotService,
|
||||
DocumentMapper documentMapper) {
|
||||
this.service = service;
|
||||
this.cleanupMapper = cleanupMapper;
|
||||
this.chunkSnapshotService = chunkSnapshotService;
|
||||
this.csvTableSnapshotService = csvTableSnapshotService;
|
||||
this.documentMapper = documentMapper;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,148 @@
|
||||
package tech.easyflow.ai.documentimport.task;
|
||||
|
||||
import org.springframework.web.multipart.MultipartFile;
|
||||
import tech.easyflow.common.filestorage.FileStorageService;
|
||||
|
||||
import java.io.ByteArrayInputStream;
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.LinkedHashSet;
|
||||
import java.util.Map;
|
||||
import java.util.Set;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
|
||||
/**
|
||||
* 文档导入测试使用的内存文件存储。
|
||||
*/
|
||||
final class InMemoryFileStorageService implements FileStorageService {
|
||||
|
||||
private final AtomicInteger sequence = new AtomicInteger();
|
||||
private final Map<String, byte[]> files = new LinkedHashMap<String, byte[]>();
|
||||
private String failingDeletePath;
|
||||
private int remainingDeleteFailures;
|
||||
|
||||
/**
|
||||
* 保存到默认测试目录。
|
||||
*
|
||||
* @param file 上传文件
|
||||
* @return 测试路径
|
||||
*/
|
||||
@Override
|
||||
public String save(MultipartFile file) {
|
||||
return save(file, "default");
|
||||
}
|
||||
|
||||
/**
|
||||
* 保存到指定测试目录。
|
||||
*
|
||||
* @param file 上传文件
|
||||
* @param prePath 测试目录
|
||||
* @return 测试路径
|
||||
*/
|
||||
@Override
|
||||
public String save(MultipartFile file, String prePath) {
|
||||
try {
|
||||
String path = prePath + "/" + sequence.incrementAndGet()
|
||||
+ "-" + file.getOriginalFilename();
|
||||
files.put(path, file.getBytes());
|
||||
return path;
|
||||
} catch (IOException error) {
|
||||
throw new IllegalStateException("测试文件保存失败", error);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 删除测试对象。
|
||||
*
|
||||
* @param path 测试路径
|
||||
*/
|
||||
@Override
|
||||
public void delete(String path) {
|
||||
if (remainingDeleteFailures > 0
|
||||
&& path != null
|
||||
&& path.equals(failingDeletePath)) {
|
||||
remainingDeleteFailures--;
|
||||
throw new IllegalStateException(
|
||||
"模拟删除失败: " + path);
|
||||
}
|
||||
files.remove(path);
|
||||
}
|
||||
|
||||
/**
|
||||
* 打开测试对象。
|
||||
*
|
||||
* @param path 测试路径
|
||||
* @return 输入流
|
||||
* @throws IOException 对象不存在
|
||||
*/
|
||||
@Override
|
||||
public InputStream readStream(String path) throws IOException {
|
||||
byte[] bytes = files.get(path);
|
||||
if (bytes == null) {
|
||||
throw new IOException("测试对象不存在: " + path);
|
||||
}
|
||||
return new ByteArrayInputStream(bytes);
|
||||
}
|
||||
|
||||
/**
|
||||
* 返回测试对象大小。
|
||||
*
|
||||
* @param path 测试路径
|
||||
* @return 字节数
|
||||
*/
|
||||
@Override
|
||||
public long getFileSize(String path) {
|
||||
byte[] bytes = files.get(path);
|
||||
return bytes == null ? -1L : bytes.length;
|
||||
}
|
||||
|
||||
/**
|
||||
* 直接放入一个源文件。
|
||||
*
|
||||
* @param path 测试路径
|
||||
* @param bytes 文件内容
|
||||
*/
|
||||
void put(String path, byte[] bytes) {
|
||||
files.put(path, bytes);
|
||||
}
|
||||
|
||||
/**
|
||||
* 返回当前对象数量。
|
||||
*
|
||||
* @return 对象数量
|
||||
*/
|
||||
int size() {
|
||||
return files.size();
|
||||
}
|
||||
|
||||
/**
|
||||
* 判断测试对象是否存在。
|
||||
*
|
||||
* @param path 测试路径
|
||||
* @return 是否存在
|
||||
*/
|
||||
boolean contains(String path) {
|
||||
return files.containsKey(path);
|
||||
}
|
||||
|
||||
/**
|
||||
* 返回当前全部测试路径快照。
|
||||
*
|
||||
* @return 路径集合
|
||||
*/
|
||||
Set<String> paths() {
|
||||
return new LinkedHashSet<String>(files.keySet());
|
||||
}
|
||||
|
||||
/**
|
||||
* 配置指定路径接下来若干次删除失败。
|
||||
*
|
||||
* @param path 测试路径
|
||||
* @param failureCount 失败次数
|
||||
*/
|
||||
void failDelete(String path, int failureCount) {
|
||||
failingDeletePath = path;
|
||||
remainingDeleteFailures = Math.max(0, failureCount);
|
||||
}
|
||||
}
|
||||
@@ -24,6 +24,7 @@ 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.DocumentImportDtos;
|
||||
import tech.easyflow.ai.documentimport.DocumentImportKeys;
|
||||
import tech.easyflow.ai.entity.DocumentChunk;
|
||||
import tech.easyflow.ai.entity.DocumentCollection;
|
||||
@@ -68,6 +69,46 @@ import java.util.concurrent.atomic.AtomicReference;
|
||||
*/
|
||||
public class KnowledgeDocumentImportTaskAppServiceTest {
|
||||
|
||||
/**
|
||||
* 验证预览翻页直接读取分片快照，避免恢复完整会话。
|
||||
*
|
||||
* @throws Exception 反射调用失败
|
||||
*/
|
||||
@Test
|
||||
public void loadPreviewPageShouldReadOnlyRequestedSnapshotRange()
|
||||
throws Exception {
|
||||
KnowledgeDocumentImportTaskAppService service =
|
||||
new KnowledgeDocumentImportTaskAppService();
|
||||
DocumentImportChunkSnapshotService snapshotService =
|
||||
Mockito.mock(DocumentImportChunkSnapshotService.class);
|
||||
setField(
|
||||
service,
|
||||
"documentImportChunkSnapshotService",
|
||||
snapshotService);
|
||||
|
||||
DocumentChunk chunk = new DocumentChunk();
|
||||
chunk.setId(BigInteger.valueOf(9001));
|
||||
Mockito.when(snapshotService.loadPage("snapshot.json", 20, 20))
|
||||
.thenReturn(List.of(chunk));
|
||||
DocumentImportDtos.PreviewSession session =
|
||||
new DocumentImportDtos.PreviewSession();
|
||||
session.setChunkSnapshotPath("snapshot.json");
|
||||
|
||||
Method method = KnowledgeDocumentImportTaskAppService.class
|
||||
.getDeclaredMethod(
|
||||
"loadPreviewPage",
|
||||
DocumentImportDtos.PreviewSession.class,
|
||||
int.class,
|
||||
int.class);
|
||||
method.setAccessible(true);
|
||||
@SuppressWarnings("unchecked")
|
||||
List<DocumentChunk> result = (List<DocumentChunk>) method.invoke(
|
||||
service, session, 2, 20);
|
||||
|
||||
Assert.assertEquals(List.of(chunk), result);
|
||||
Mockito.verify(snapshotService).loadPage("snapshot.json", 20, 20);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证待处理任务重新投递只更新必要字段，避免自定义查询结果覆盖非空列。
|
||||
*
|
||||
|
||||
Reference in New Issue
Block a user