package com.bruce.rag; import com.bruce.common.config.AttachmentProperties; import com.bruce.common.document.parse.DocumentParseContext; import com.bruce.common.document.parse.DocumentParseResult; import com.bruce.common.document.parse.DocumentParser; import com.bruce.common.document.parse.DocumentParserFactory; import com.bruce.common.domain.entity.SysAttachment; import com.bruce.common.service.ISysAttachmentService; import com.bruce.rag.dto.request.RagDocumentParseRequest; import com.bruce.rag.dto.response.RagDocumentParseResponse; import com.bruce.rag.entity.RagDocument; import com.bruce.rag.mapper.RagDocumentMapper; import com.bruce.rag.service.IRagDocumentParseResultService; import com.bruce.rag.enums.RagParseStatusEnum; import com.bruce.rag.service.impl.RagDocumentParseServiceImpl; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; import java.nio.file.Files; import java.nio.file.Path; import java.util.List; import java.util.Map; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.times; import static org.mockito.Mockito.when; import static org.mockito.Mockito.mock; @ExtendWith(MockitoExtension.class) class RagDocumentParseServiceImplTests { @TempDir private Path tempDir; @Mock private RagDocumentMapper ragDocumentMapper; @Mock private ISysAttachmentService sysAttachmentService; @Test void parseShouldUpdateStatusAndReturnParseResponse() throws Exception { Path file = tempDir.resolve("rag").resolve("people.txt"); Files.createDirectories(file.getParent()); Files.writeString(file, "people profiles"); RagDocument document = new RagDocument(); document.setId(1001L); document.setStoreId(2002L); document.setAttachmentId(3003L); document.setParseStatus(RagParseStatusEnum.UPLOADED.name()); SysAttachment attachment = new SysAttachment(); attachment.setId(3003L); attachment.setOriginalName("people.txt"); attachment.setFileSuffix("txt"); attachment.setContentType("text/plain"); attachment.setFilePath("rag/people.txt"); AttachmentProperties attachmentProperties = new AttachmentProperties(); attachmentProperties.setBasePath(tempDir.toString()); DocumentParser parser = new FixedDocumentParser("people profiles"); RagDocumentParseServiceImpl service = new RagDocumentParseServiceImpl( ragDocumentMapper, sysAttachmentService, attachmentProperties, new DocumentParserFactory(List.of(parser)), mock(IRagDocumentParseResultService.class) ); when(ragDocumentMapper.selectById(1001L)).thenReturn(document); when(sysAttachmentService.getById(3003L)).thenReturn(attachment); when(ragDocumentMapper.updateById(any(RagDocument.class))).thenReturn(1); RagDocumentParseResponse response = service.parse(1001L); assertEquals(1001L, response.getDocumentId()); assertEquals(RagParseStatusEnum.PARSED.name(), response.getParseStatus()); assertEquals(15, response.getTextLength()); assertEquals("fixed", response.getMetadata().get("parser")); ArgumentCaptor captor = ArgumentCaptor.forClass(RagDocument.class); verify(ragDocumentMapper, times(2)).updateById(captor.capture()); List updates = captor.getAllValues(); assertEquals(RagParseStatusEnum.PARSING.name(), updates.get(0).getParseStatus()); assertEquals(RagParseStatusEnum.PARSED.name(), updates.get(1).getParseStatus()); assertTrue(parser.supports(new DocumentParseContext())); } @Test void parseShouldSupportBatchRequest() throws Exception { Path file = tempDir.resolve("rag").resolve("batch.txt"); Files.createDirectories(file.getParent()); Files.writeString(file, "batch profiles"); RagDocument document = new RagDocument(); document.setId(1002L); document.setStoreId(2002L); document.setAttachmentId(3004L); document.setParseStatus(RagParseStatusEnum.UPLOADED.name()); SysAttachment attachment = new SysAttachment(); attachment.setId(3004L); attachment.setOriginalName("batch.txt"); attachment.setFileSuffix("txt"); attachment.setContentType("text/plain"); attachment.setFilePath("rag/batch.txt"); AttachmentProperties attachmentProperties = new AttachmentProperties(); attachmentProperties.setBasePath(tempDir.toString()); RagDocumentParseServiceImpl service = new RagDocumentParseServiceImpl( ragDocumentMapper, sysAttachmentService, attachmentProperties, new DocumentParserFactory(List.of(new FixedDocumentParser("batch profiles"))), mock(IRagDocumentParseResultService.class) ); RagDocumentParseRequest request = new RagDocumentParseRequest(); request.setDocumentIds(List.of(1002L)); when(ragDocumentMapper.selectById(1002L)).thenReturn(document); when(sysAttachmentService.getById(3004L)).thenReturn(attachment); when(ragDocumentMapper.updateById(any(RagDocument.class))).thenReturn(1); List responses = service.parse(request); assertEquals(1, responses.size()); assertEquals(1002L, responses.getFirst().getDocumentId()); assertEquals(RagParseStatusEnum.PARSED.name(), responses.getFirst().getParseStatus()); } @Test void parseShouldRejectEmptyDocumentIds() { AttachmentProperties attachmentProperties = new AttachmentProperties(); attachmentProperties.setBasePath(tempDir.toString()); RagDocumentParseServiceImpl service = new RagDocumentParseServiceImpl( ragDocumentMapper, sysAttachmentService, attachmentProperties, new DocumentParserFactory(List.of(new FixedDocumentParser("batch profiles"))), mock(IRagDocumentParseResultService.class) ); RagDocumentParseRequest request = new RagDocumentParseRequest(); request.setDocumentIds(List.of()); assertThrows(IllegalArgumentException.class, () -> service.parse(request)); } private static class FixedDocumentParser implements DocumentParser { private final String text; private FixedDocumentParser(String text) { this.text = text; } @Override public boolean supports(DocumentParseContext context) { return true; } @Override public DocumentParseResult parse(DocumentParseContext context) { DocumentParseResult result = new DocumentParseResult(); result.setText(text); result.setTextLength(text.length()); result.setMetadata(Map.of("parser", "fixed")); return result; } } }