package com.zy.ai.service;
|
|
import com.alibaba.fastjson.JSONObject;
|
import com.baomidou.mybatisplus.core.conditions.Wrapper;
|
import com.zy.ai.domain.autotune.AutoTuneApplyRequest;
|
import com.zy.ai.domain.autotune.AutoTuneApplyResult;
|
import com.zy.ai.domain.autotune.AutoTuneChangeCommand;
|
import com.zy.ai.domain.autotune.AutoTuneSnapshot;
|
import com.zy.ai.entity.AiAutoTuneChange;
|
import com.zy.ai.entity.AiAutoTuneJob;
|
import com.zy.ai.entity.AiPromptTemplate;
|
import com.zy.ai.entity.ChatCompletionRequest;
|
import com.zy.ai.entity.ChatCompletionResponse;
|
import com.zy.ai.mcp.service.SpringAiMcpToolManager;
|
import com.zy.ai.mcp.tool.AutoTuneMcpTools;
|
import com.zy.ai.service.impl.AutoTuneAgentServiceImpl;
|
import org.junit.jupiter.api.BeforeEach;
|
import org.junit.jupiter.api.Test;
|
import org.junit.jupiter.api.extension.ExtendWith;
|
import org.mockito.ArgumentCaptor;
|
import org.mockito.Mock;
|
import org.mockito.junit.jupiter.MockitoExtension;
|
|
import java.util.ArrayList;
|
import java.util.Collections;
|
import java.util.List;
|
import java.util.Map;
|
|
import static org.junit.jupiter.api.Assertions.assertEquals;
|
import static org.junit.jupiter.api.Assertions.assertFalse;
|
import static org.junit.jupiter.api.Assertions.assertSame;
|
import static org.junit.jupiter.api.Assertions.assertTrue;
|
import static org.mockito.ArgumentMatchers.any;
|
import static org.mockito.ArgumentMatchers.anyDouble;
|
import static org.mockito.ArgumentMatchers.anyInt;
|
import static org.mockito.Mockito.verify;
|
import static org.mockito.Mockito.when;
|
|
@ExtendWith(MockitoExtension.class)
|
class AutoTuneCoordinatorServiceImplTest {
|
|
private AutoTuneMcpTools tools;
|
|
@Mock
|
private AutoTuneSnapshotService autoTuneSnapshotService;
|
@Mock
|
private AutoTuneApplyService autoTuneApplyService;
|
@Mock
|
private AiAutoTuneJobService aiAutoTuneJobService;
|
@Mock
|
private AiAutoTuneChangeService aiAutoTuneChangeService;
|
@Mock
|
private LlmChatService llmChatService;
|
@Mock
|
private SpringAiMcpToolManager mcpToolManager;
|
@Mock
|
private AiPromptTemplateService aiPromptTemplateService;
|
|
@BeforeEach
|
void setUp() {
|
tools = new AutoTuneMcpTools(
|
autoTuneSnapshotService,
|
autoTuneApplyService,
|
aiAutoTuneJobService,
|
aiAutoTuneChangeService);
|
}
|
|
@Test
|
void snapshotToolDelegatesToSnapshotService() {
|
AutoTuneSnapshot snapshot = new AutoTuneSnapshot();
|
when(autoTuneSnapshotService.buildSnapshot()).thenReturn(snapshot);
|
|
AutoTuneSnapshot result = tools.getAutoTuneSnapshot();
|
|
assertSame(snapshot, result);
|
verify(autoTuneSnapshotService).buildSnapshot();
|
}
|
|
@Test
|
void recentJobsReturnsBoundedCompactSummariesWithChanges() {
|
AiAutoTuneJob job = new AiAutoTuneJob();
|
job.setId(7L);
|
job.setTriggerType("agent");
|
job.setStatus("success");
|
job.setSummary("applied");
|
job.setSuccessCount(1);
|
job.setRejectCount(0);
|
AiAutoTuneChange change = new AiAutoTuneChange();
|
change.setJobId(7L);
|
change.setTargetType("sys_config");
|
change.setTargetKey("conveyorStationTaskLimit");
|
change.setRequestedValue("12");
|
change.setResultStatus("success");
|
|
when(aiAutoTuneJobService.list(any(Wrapper.class))).thenReturn(Collections.singletonList(job));
|
when(aiAutoTuneChangeService.list(any(Wrapper.class))).thenReturn(Collections.singletonList(change));
|
|
List<Map<String, Object>> result = tools.getRecentAutoTuneJobs(99);
|
|
assertEquals(1, result.size());
|
assertEquals(7L, result.get(0).get("id"));
|
assertFalse(result.get(0).containsKey("reasoningDigest"));
|
List<?> changes = (List<?>) result.get(0).get("changes");
|
assertEquals(1, changes.size());
|
|
ArgumentCaptor<Wrapper<AiAutoTuneJob>> wrapperCaptor = ArgumentCaptor.forClass(Wrapper.class);
|
verify(aiAutoTuneJobService).list(wrapperCaptor.capture());
|
assertTrue(wrapperCaptor.getValue().getSqlSegment().contains("limit 20"));
|
}
|
|
@Test
|
void applyToolDelegatesToApplyServiceWithDryRunAndChanges() {
|
AutoTuneApplyResult expected = new AutoTuneApplyResult();
|
expected.setDryRun(true);
|
when(autoTuneApplyService.apply(any(AutoTuneApplyRequest.class))).thenReturn(expected);
|
|
AutoTuneChangeCommand command = new AutoTuneChangeCommand();
|
command.setTargetType("sys_config");
|
command.setTargetKey("conveyorStationTaskLimit");
|
command.setNewValue("12");
|
List<AutoTuneChangeCommand> changes = Collections.singletonList(command);
|
|
AutoTuneApplyResult result = tools.applyAutoTuneChanges("reduce congestion", 10, "agent", true, changes);
|
|
assertSame(expected, result);
|
ArgumentCaptor<AutoTuneApplyRequest> captor = ArgumentCaptor.forClass(AutoTuneApplyRequest.class);
|
verify(autoTuneApplyService).apply(captor.capture());
|
assertEquals("reduce congestion", captor.getValue().getReason());
|
assertEquals(10, captor.getValue().getAnalysisIntervalMinutes());
|
assertEquals("agent", captor.getValue().getTriggerType());
|
assertEquals(Boolean.TRUE, captor.getValue().getDryRun());
|
assertSame(changes, captor.getValue().getChanges());
|
}
|
|
@Test
|
void rollbackToolDelegatesToApplyServiceRollback() {
|
AutoTuneApplyResult expected = new AutoTuneApplyResult();
|
when(autoTuneApplyService.rollbackLastSuccessfulJob("bad result")).thenReturn(expected);
|
|
AutoTuneApplyResult result = tools.revertLastAutoTuneJob("bad result");
|
|
assertSame(expected, result);
|
verify(autoTuneApplyService).rollbackLastSuccessfulJob("bad result");
|
}
|
|
@Test
|
void agentExecutesSnapshotDryRunAndRealApplyToolSequence() {
|
AutoTuneAgentServiceImpl service = new AutoTuneAgentServiceImpl(
|
llmChatService,
|
mcpToolManager,
|
aiPromptTemplateService);
|
AiPromptTemplate promptTemplate = new AiPromptTemplate();
|
promptTemplate.setContent("system prompt");
|
when(aiPromptTemplateService.resolvePublished("wcs_auto_tune_dispatch")).thenReturn(promptTemplate);
|
when(mcpToolManager.buildOpenAiTools()).thenReturn(Collections.singletonList(Collections.singletonMap("type", "function")));
|
when(mcpToolManager.callTool(any(), any(JSONObject.class))).thenReturn(Collections.singletonMap("ok", true));
|
when(llmChatService.chatCompletion(any(), anyDouble(), anyInt(), any()))
|
.thenReturn(
|
response("read snapshot", toolCall("call_1", "wcs_local_dispatch_get_auto_tune_snapshot", "{}"), 10, 5),
|
response("dry-run first", toolCall("call_2", "wcs_local_dispatch_apply_auto_tune_changes",
|
"{\"dryRun\":true,\"changes\":[{\"targetType\":\"sys_config\",\"targetKey\":\"conveyorStationTaskLimit\",\"newValue\":\"12\"}]}"), 11, 6),
|
response("apply after dry-run", toolCall("call_3", "wcs_local_dispatch_apply_auto_tune_changes",
|
"{\"dryRun\":false,\"changes\":[{\"targetType\":\"sys_config\",\"targetKey\":\"conveyorStationTaskLimit\",\"newValue\":\"12\"}]}"), 12, 7),
|
response("已完成自动调参", null, 13, 8)
|
);
|
|
AutoTuneAgentService.AutoTuneAgentResult result = service.runAutoTune("scheduler");
|
|
assertTrue(result.getSuccess());
|
assertEquals("scheduler", result.getTriggerType());
|
assertEquals(3, result.getToolCallCount());
|
assertEquals(4, result.getLlmCallCount());
|
assertEquals(46L, result.getPromptTokens());
|
assertEquals(26L, result.getCompletionTokens());
|
assertEquals(72L, result.getTotalTokens());
|
assertTrue(result.getSummary().contains("已完成自动调参"));
|
|
ArgumentCaptor<String> toolNameCaptor = ArgumentCaptor.forClass(String.class);
|
ArgumentCaptor<JSONObject> argumentCaptor = ArgumentCaptor.forClass(JSONObject.class);
|
verify(mcpToolManager, org.mockito.Mockito.times(3)).callTool(toolNameCaptor.capture(), argumentCaptor.capture());
|
assertEquals("wcs_local_dispatch_get_auto_tune_snapshot", toolNameCaptor.getAllValues().get(0));
|
assertEquals("wcs_local_dispatch_apply_auto_tune_changes", toolNameCaptor.getAllValues().get(1));
|
assertEquals(Boolean.TRUE, argumentCaptor.getAllValues().get(1).getBoolean("dryRun"));
|
assertEquals(Boolean.FALSE, argumentCaptor.getAllValues().get(2).getBoolean("dryRun"));
|
}
|
|
private ChatCompletionResponse response(String content,
|
ChatCompletionRequest.ToolCall toolCall,
|
int promptTokens,
|
int completionTokens) {
|
ChatCompletionRequest.Message message = new ChatCompletionRequest.Message();
|
message.setRole("assistant");
|
message.setContent(content);
|
if (toolCall != null) {
|
message.setTool_calls(Collections.singletonList(toolCall));
|
}
|
|
ChatCompletionResponse.Choice choice = new ChatCompletionResponse.Choice();
|
choice.setIndex(0);
|
choice.setMessage(message);
|
|
ChatCompletionResponse.Usage usage = new ChatCompletionResponse.Usage();
|
usage.setPromptTokens(promptTokens);
|
usage.setCompletionTokens(completionTokens);
|
usage.setTotalTokens(promptTokens + completionTokens);
|
|
ChatCompletionResponse response = new ChatCompletionResponse();
|
List<ChatCompletionResponse.Choice> choices = new ArrayList<>();
|
choices.add(choice);
|
response.setChoices(choices);
|
response.setUsage(usage);
|
return response;
|
}
|
|
private ChatCompletionRequest.ToolCall toolCall(String id, String name, String arguments) {
|
ChatCompletionRequest.Function function = new ChatCompletionRequest.Function();
|
function.setName(name);
|
function.setArguments(arguments);
|
|
ChatCompletionRequest.ToolCall toolCall = new ChatCompletionRequest.ToolCall();
|
toolCall.setId(id);
|
toolCall.setType("function");
|
toolCall.setFunction(function);
|
return toolCall;
|
}
|
}
|
