common_agent/common-agent-rag/src/test/java/com/bruce/rag/RagDocumentParseServiceImplTests.java

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<RagDocument> captor = ArgumentCaptor.forClass(RagDocument.class);
        verify(ragDocumentMapper, times(2)).updateById(captor.capture());
        List<RagDocument> 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<RagDocumentParseResponse> 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;
        }
    }
}