SuperBizAgent-java/src/test/java/com/superbiz/agent/service/DiagnosisTraceServiceTest.java

package com.superbiz.agent.service;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.superbiz.agent.domain.entity.AgentStep;
import com.superbiz.agent.domain.entity.DiagnosisSession;
import com.superbiz.agent.domain.entity.ToolInvocation;
import com.superbiz.agent.dto.DiagnosisTraceResponse;
import com.superbiz.agent.exception.SessionNotFoundException;
import com.superbiz.agent.repository.AgentStepRepository;
import com.superbiz.agent.repository.DiagnosisSessionRepository;
import com.superbiz.agent.repository.ToolInvocationRepository;
import org.junit.jupiter.api.Test;

import java.time.LocalDateTime;
import java.util.List;
import java.util.Optional;

import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.Mockito.*;

class DiagnosisTraceServiceTest {

    private final DiagnosisSessionRepository diagnosisSessionRepository = mock(DiagnosisSessionRepository.class);
    private final AgentStepRepository agentStepRepository = mock(AgentStepRepository.class);
    private final ToolInvocationRepository toolInvocationRepository = mock(ToolInvocationRepository.class);
    private final DiagnosisTraceService service = new DiagnosisTraceService(
            diagnosisSessionRepository,
            agentStepRepository,
            toolInvocationRepository,
            new ObjectMapper()
    );

    @Test
    void getTraceAggregatesSessionStepsAndTools() {
        String sessionId = "trace-session-001";
        LocalDateTime now = LocalDateTime.of(2026, 7, 3, 14, 30);
        DiagnosisSession session = DiagnosisSession.builder()
                .id(1L)
                .sessionId(sessionId)
                .query("payment timeout")
                .status("SUCCESS")
                .agentFlow("COMPLEX")
                .totalDurationMs(1200)
                .totalTokenCount(300)
                .stepCount(2)
                .toolCallCount(1)
                .answer("restart payment gateway pool")
                .selfEvaluation("{\"verifier_evaluation\":{\"verdict\":\"PASS\"},\"aiops_rule_evaluation\":{\"verdict\":\"WARN\"}}")
                .feedback("useful")
                .createdAt(now)
                .updatedAt(now)
                .build();
        AgentStep step = AgentStep.builder()
                .id(10L)
                .sessionId(sessionId)
                .stepIndex(1)
                .agentName("chat_executor")
                .modelInput("input")
                .modelOutput("output")
                .thought("executor finished")
                .hasToolCall(true)
                .durationMs(500)
                .tokenCount(100)
                .createdAt(now)
                .build();
        ToolInvocation invocation = ToolInvocation.builder()
                .id(20L)
                .sessionId(sessionId)
                .stepId(10L)
                .toolName("lookup_knowledge")
                .inputParams("{\"query\":\"ERR_TIMEOUT\"}")
                .outputPreview("payment timeout doc")
                .outputLength(19)
                .retrievalLayer("L0")
                .l0MatchCount(1)
                .l1MatchCount(0)
                .isTruncated(false)
                .relevanceLevel("HIGHLY_RELEVANT")
                .dedupReason("FIRST_HIT")
                .retrievalDetails("{\"documents\":[\"payment-errors.md\"]}")
                .durationMs(80)
                .success(true)
                .createdAt(now)
                .build();

        when(diagnosisSessionRepository.findBySessionId(sessionId)).thenReturn(Optional.of(session));
        when(agentStepRepository.findBySessionId(sessionId)).thenReturn(List.of(step));
        when(toolInvocationRepository.findBySessionIdOrderByIdAsc(sessionId)).thenReturn(List.of(invocation));

        DiagnosisTraceResponse response = service.getTrace(sessionId);

        assertEquals(sessionId, response.getSession().getSessionId());
        assertEquals("payment timeout", response.getSession().getQuery());
        assertEquals("PASS", ((java.util.Map<?, ?>) response.getSession()
                .getSelfEvaluation()
                .get("verifier_evaluation")).get("verdict"));
        assertEquals(1, response.getSteps().size());
        assertEquals("chat_executor", response.getSteps().get(0).getAgentName());
        assertEquals(1, response.getToolInvocations().size());
        assertEquals("ERR_TIMEOUT", response.getToolInvocations().get(0).getInputParams().get("query"));
        assertEquals(2, response.getSummary().getPersistedStepCount());
        assertEquals(1, response.getSummary().getReturnedStepCount());
        assertEquals(1, response.getSummary().getPersistedToolCallCount());
        assertEquals(1, response.getSummary().getReturnedToolCallCount());
        assertTrue(response.getSummary().isHasVerifierEvaluation());
        assertTrue(response.getSummary().isHasAiOpsRuleEvaluation());
        assertTrue(response.getSummary().isHasFeedback());
    }

    @Test
    void getTraceOrdersStepsByCreationTimeAndIdNotPerAgentStepIndex() {
        String sessionId = "trace-session-ordered";
        LocalDateTime base = LocalDateTime.of(2026, 7, 7, 0, 0);
        DiagnosisSession session = DiagnosisSession.builder()
                .id(1L)
                .sessionId(sessionId)
                .query("mysql pool issue")
                .status("SUCCESS")
                .stepCount(4)
                .toolCallCount(0)
                .createdAt(base)
                .updatedAt(base)
                .build();
        AgentStep planner = step(1L, sessionId, 0, "planner", base.plusSeconds(1));
        AgentStep executor0 = step(2L, sessionId, 0, "executor", base.plusSeconds(2));
        AgentStep executor1 = step(3L, sessionId, 1, "executor", base.plusSeconds(3));
        AgentStep verifier = step(4L, sessionId, 0, "verifier", base.plusSeconds(4));

        when(diagnosisSessionRepository.findBySessionId(sessionId)).thenReturn(Optional.of(session));
        when(agentStepRepository.findBySessionId(sessionId))
                .thenReturn(List.of(planner, executor0, verifier, executor1));
        when(toolInvocationRepository.findBySessionIdOrderByIdAsc(sessionId)).thenReturn(List.of());

        DiagnosisTraceResponse response = service.getTrace(sessionId);

        assertEquals(List.of("planner", "executor", "executor", "verifier"),
                response.getSteps().stream().map(DiagnosisTraceResponse.AgentStepTrace::getAgentName).toList());
        assertEquals(List.of(0, 0, 1, 0),
                response.getSteps().stream().map(DiagnosisTraceResponse.AgentStepTrace::getStepIndex).toList());
    }

    @Test
    void getTraceThrowsWhenSessionMissing() {
        String sessionId = "missing-session";
        when(diagnosisSessionRepository.findBySessionId(sessionId)).thenReturn(Optional.empty());

        assertThrows(SessionNotFoundException.class, () -> service.getTrace(sessionId));

        verify(diagnosisSessionRepository).findBySessionId(sessionId);
        verifyNoInteractions(agentStepRepository, toolInvocationRepository);
    }

    private AgentStep step(Long id, String sessionId, int stepIndex, String agentName, LocalDateTime createdAt) {
        return AgentStep.builder()
                .id(id)
                .sessionId(sessionId)
                .stepIndex(stepIndex)
                .agentName(agentName)
                .createdAt(createdAt)
                .build();
    }
}