BJ_ZTEITS-Developers / parking-cloud-yudao

Commit 2c781729f2da441bb3825b31515b393c8e19f8d8

Authored by YunaiV
1 parent 8d94c7c1

fix(ai): 修复 Midjourney 操作、模型同步及代理请求

Inline Side-by-side
Showing 4 changed files with 654 additions and 8 deletions
yudao-module-ai/src/main/java/cn/iocoder/yudao/module/ai/framework/ai/core/model/midjourney/api/MidjourneyApi.java
  View file @2c78172
... ... @@ -33,7 +33,7 @@ public class MidjourneyApi {
33 33 private final Predicate<HttpStatusCode> STATUS_PREDICATE = status -> !status.is2xxSuccessful();
34 34  
35 35 private final Function<Object, Function<ClientResponse, Mono<? extends Throwable>>> EXCEPTION_FUNCTION =
36   - reqParam -> response -> response.bodyToMono(String.class).handle((responseBody, sink) -> {
  36 + reqParam -> response -> response.bodyToMono(String.class).defaultIfEmpty("").handle((responseBody, sink) -> {
37 37 HttpRequest request = response.request();
38 38 log.error("[midjourney-api] 调用失败！请求方式:[{}]，请求地址:[{}]，请求参数:[{}]，响应数据: [{}]",
39 39 request.getMethod(), request.getURI(), reqParam, responseBody);
... ... @@ -53,6 +53,10 @@ public class MidjourneyApi {
53 53 .defaultHeaders(httpHeaders -> {
54 54 httpHeaders.setContentType(MediaType.APPLICATION_JSON);
55 55 httpHeaders.setBearerAuth(apiKey);
  56 + // 兼容 novicezk/midjourney-proxy：开启 mj.api-secret 后，通过 mj-api-secret 请求头鉴权
  57 + if (StrUtil.isNotEmpty(apiKey)) {
  58 + httpHeaders.set("mj-api-secret", apiKey);
  59 + }
56 60 })
57 61 .build();
58 62 this.notifyUrl = notifyUrl;
... ...
yudao-module-ai/src/main/java/cn/iocoder/yudao/module/ai/service/image/AiImageServiceImpl.java
  View file @2c78172
... ... @@ -46,7 +46,9 @@ import java.util.Map;
46 46  
47 47 import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
48 48 import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertMap;
  49 +import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertMultiMap;
49 50 import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertSet;
  51 +import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.filterList;
50 52 import static cn.iocoder.yudao.module.ai.enums.ErrorCodeConstants.*;
51 53  
52 54 /**
... ... @@ -229,10 +231,11 @@ public class AiImageServiceImpl implements AiImageService {
229 231 // 3. 调用 Midjourney Proxy 提交任务
230 232 List<String> base64Array = StrUtil.isBlank(drawReqVO.getReferImageUrl()) ? null :
231 233 Collections.singletonList("data:image/jpeg;base64,".concat(Base64.encode(HttpUtil.downloadBytes(drawReqVO.getReferImageUrl()))));
  234 + // 注意：绘图参数（--ar、--v 等）需要拼接到 prompt 中提交；Midjourney Proxy 的 state 只是自定义参数，不会提交给 Midjourney
  235 + String state = MidjourneyApi.ImagineRequest.buildState(drawReqVO.getWidth(),
  236 + drawReqVO.getHeight(), drawReqVO.getVersion(), model.getModel());
232 237 MidjourneyApi.ImagineRequest imagineRequest = new MidjourneyApi.ImagineRequest(
233   - base64Array, drawReqVO.getPrompt(),null,
234   - MidjourneyApi.ImagineRequest.buildState(drawReqVO.getWidth(),
235   - drawReqVO.getHeight(), drawReqVO.getVersion(), model.getModel()));
  238 + base64Array, drawReqVO.getPrompt().concat(state), null, state);
236 239 MidjourneyApi.SubmitResponse imagineResponse = midjourneyApi.imagine(imagineRequest);
237 240  
238 241 // 4.1 情况一【失败】：抛出业务异常
... ... @@ -256,8 +259,25 @@ public class AiImageServiceImpl implements AiImageService {
256 259 if (CollUtil.isEmpty(images)) {
257 260 return 0;
258 261 }
259   - // 1.2 调用 Midjourney Proxy 获取任务进展
260   - MidjourneyApi midjourneyApi = modelService.getMidjourneyApi(images.get(0).getModelId());
  262 + // 1.2 特殊：未关联模型的 image（历史 Action 生成的），无法确定 Midjourney Proxy，跳过同步，依赖回调更新
  263 + List<AiImageDO> unboundImages = filterList(images, image -> image.getModelId() == null);
  264 + if (CollUtil.isNotEmpty(unboundImages)) {
  265 + log.warn("[midjourneySync][image({}) 未关联模型，跳过同步]", convertSet(unboundImages, AiImageDO::getId));
  266 + }
  267 +
  268 + // 2. 按照模型分组同步：不同模型可能对应不同的 Midjourney Proxy，需要使用各自的 API 查询
  269 + int count = 0;
  270 + Map<Long, List<AiImageDO>> modelImagesMap = convertMultiMap(
  271 + filterList(images, image -> image.getModelId() != null), AiImageDO::getModelId);
  272 + for (Map.Entry<Long, List<AiImageDO>> entry : modelImagesMap.entrySet()) {
  273 + count += midjourneySync(entry.getKey(), entry.getValue());
  274 + }
  275 + return count;
  276 + }
  277 +
  278 + private int midjourneySync(Long modelId, List<AiImageDO> images) {
  279 + // 1. 调用 Midjourney Proxy 获取任务进展
  280 + MidjourneyApi midjourneyApi = modelService.getMidjourneyApi(modelId);
261 281 List<MidjourneyApi.Notify> taskList = midjourneyApi.getTaskList(convertSet(images, AiImageDO::getTaskId));
262 282 Map<String, MidjourneyApi.Notify> taskMap = convertMap(taskList, MidjourneyApi.Notify::id);
263 283  
... ... @@ -336,7 +356,7 @@ public class AiImageServiceImpl implements AiImageService {
336 356  
337 357 // 2. 调用 Midjourney Proxy 提交任务
338 358 MidjourneyApi.SubmitResponse actionResponse = midjourneyApi.action(
339   - new MidjourneyApi.ActionRequest(button.customId(), image.getTaskId(), null));
  359 + new MidjourneyApi.ActionRequest(image.getTaskId(), button.customId(), null));
340 360 if (!MidjourneyApi.SubmitCodeEnum.SUCCESS_CODES.contains(actionResponse.code())) {
341 361 String description = actionResponse.description().contains("quota_not_enough") ?
342 362 "账户余额不足" : actionResponse.description();
... ... @@ -347,7 +367,7 @@ public class AiImageServiceImpl implements AiImageService {
347 367 AiImageDO newImage = new AiImageDO().setUserId(image.getUserId()).setPublicStatus(false).setPrompt(image.getPrompt())
348 368 .setStatus(AiImageStatusEnum.IN_PROGRESS.getStatus())
349 369 .setPlatform(AiPlatformEnum.MIDJOURNEY.getPlatform())
350   - .setModel(image.getModel()).setWidth(image.getWidth()).setHeight(image.getHeight())
  370 + .setModelId(image.getModelId()).setModel(image.getModel()).setWidth(image.getWidth()).setHeight(image.getHeight())
351 371 .setOptions(image.getOptions()).setTaskId(actionResponse.result());
352 372 imageMapper.insert(newImage);
353 373 return newImage.getId();
... ...
yudao-module-ai/src/test/java/cn/iocoder/yudao/module/ai/framework/ai/core/model/midjourney/api/MidjourneyApiUnitTest.java 0 → 100644
  View file @2c78172
  1 +package cn.iocoder.yudao.module.ai.framework.ai.core.model.midjourney.api;
  2 +
  3 +import cn.iocoder.yudao.framework.test.core.ut.BaseMockitoUnitTest;
  4 +import org.junit.jupiter.api.BeforeEach;
  5 +import org.junit.jupiter.api.Test;
  6 +import org.junit.jupiter.params.ParameterizedTest;
  7 +import org.junit.jupiter.params.provider.NullAndEmptySource;
  8 +import org.junit.jupiter.params.provider.ValueSource;
  9 +import org.mockito.ArgumentCaptor;
  10 +import org.mockito.Mock;
  11 +import org.mockito.MockedStatic;
  12 +import org.springframework.http.HttpHeaders;
  13 +import org.springframework.http.HttpMethod;
  14 +import org.springframework.http.HttpStatus;
  15 +import org.springframework.http.MediaType;
  16 +import org.springframework.web.reactive.function.client.ClientRequest;
  17 +import org.springframework.web.reactive.function.client.ClientResponse;
  18 +import org.springframework.web.reactive.function.client.ExchangeFunction;
  19 +import org.springframework.web.reactive.function.client.WebClient;
  20 +import reactor.core.publisher.Mono;
  21 +
  22 +import java.net.URI;
  23 +import java.util.Collections;
  24 +import java.util.List;
  25 +
  26 +import static org.junit.jupiter.api.Assertions.*;
  27 +import static org.mockito.ArgumentMatchers.any;
  28 +import static org.mockito.Mockito.*;
  29 +
  30 +/**
  31 + * {@link MidjourneyApi} 的单元测试类
  32 + *
  33 + * 通过 mock {@link ExchangeFunction} 截获真实 {@link WebClient} 发出的 {@link ClientRequest}，不访问外部服务、不监听本地端口。
  34 + * 重点覆盖 Midjourney Proxy 的鉴权请求头：
  35 + * 1. 保留原有的 Authorization: Bearer 请求头，兼容托管服务
  36 + * 2. 补充 mj-api-secret 请求头，兼容 novicezk/midjourney-proxy 开启 mj.api-secret 后的鉴权
  37 + * 3. 非 2xx 且无响应 body 时，仍抛出异常，不被吞掉
  38 + *
  39 + * @author 芋道源码
  40 + */
  41 +public class MidjourneyApiUnitTest extends BaseMockitoUnitTest {
  42 +
  43 + private static final String BASE_URL = "https://mj-proxy.example.com/mj";
  44 + private static final String API_KEY = "test-mj-api-secret";
  45 + private static final String NOTIFY_URL = "https://admin.example.com/admin-api/ai/image/midjourney/notify";
  46 +
  47 + @Mock
  48 + private ExchangeFunction exchangeFunction;
  49 +
  50 + private MidjourneyApi midjourneyApi;
  51 +
  52 + @BeforeEach
  53 + public void setUp() {
  54 + midjourneyApi = createMidjourneyApi(API_KEY);
  55 + }
  56 +
  57 + // ========== imagine ==========
  58 +
  59 + @Test
  60 + public void testImagine_success() {
  61 + // mock 方法（Midjourney Proxy 返回）
  62 + mockResponse(HttpStatus.OK, MediaType.APPLICATION_JSON_VALUE,
  63 + "{\"code\":\"1\",\"description\":\"提交成功\",\"result\":\"task-1\"}");
  64 + // 准备参数
  65 + MidjourneyApi.ImagineRequest request = new MidjourneyApi.ImagineRequest(null, "一只猫", null, null);
  66 +
  67 + // 调用
  68 + MidjourneyApi.SubmitResponse response = midjourneyApi.imagine(request);
  69 + // 断言：返回
  70 + assertEquals("1", response.code());
  71 + assertEquals("提交成功", response.description());
  72 + assertEquals("task-1", response.result());
  73 + assertEquals(NOTIFY_URL, request.getNotifyHook());
  74 + // 断言：请求
  75 + assertAuthorizedRequest(captureRequest(), "/submit/imagine");
  76 + }
  77 +
  78 + @ParameterizedTest
  79 + @NullAndEmptySource
  80 + public void testImagine_withoutApiKey(String apiKey) {
  81 + // 准备参数（未配置 apiKey）
  82 + MidjourneyApi api = createMidjourneyApi(apiKey);
  83 + // mock 方法（Midjourney Proxy 返回）
  84 + mockResponse(HttpStatus.OK, MediaType.APPLICATION_JSON_VALUE,
  85 + "{\"code\":\"1\",\"description\":\"提交成功\",\"result\":\"task-1\"}");
  86 +
  87 + // 调用
  88 + MidjourneyApi.SubmitResponse response = api.imagine(
  89 + new MidjourneyApi.ImagineRequest(null, "一只猫", null, null));
  90 + // 断言：返回
  91 + assertEquals("task-1", response.result());
  92 + // 断言：请求（Authorization 保持原行为，不补充 mj-api-secret）
  93 + ClientRequest clientRequest = captureRequest();
  94 + assertEquals(URI.create("https://mj-proxy.example.com/mj/submit/imagine"), clientRequest.url());
  95 + assertEquals("Bearer " + apiKey, clientRequest.headers().getFirst(HttpHeaders.AUTHORIZATION));
  96 + assertFalse(clientRequest.headers().containsKey("mj-api-secret"));
  97 + }
  98 +
  99 + @Test
  100 + public void testImagine_unauthorized() {
  101 + // mock 方法（Midjourney Proxy 鉴权失败）
  102 + mockResponse(HttpStatus.UNAUTHORIZED, MediaType.TEXT_PLAIN_VALUE, "Unauthorized");
  103 +
  104 + // 调用，并断言异常
  105 + IllegalStateException exception = assertThrows(IllegalStateException.class, () -> midjourneyApi.imagine(
  106 + new MidjourneyApi.ImagineRequest(null, "一只猫", null, null)));
  107 + assertEquals("[midjourney-api] 调用失败！", exception.getMessage());
  108 + // 断言：请求
  109 + assertAuthorizedRequest(captureRequest(), "/submit/imagine");
  110 + }
  111 +
  112 + @ParameterizedTest
  113 + @ValueSource(ints = {401, 500})
  114 + public void testImagine_errorWithoutBody(int status) {
  115 + // mock 方法（Midjourney Proxy 失败，且无响应 body）
  116 + mockResponseWithoutBody(status);
  117 +
  118 + // 调用，并断言异常
  119 + IllegalStateException exception = assertThrows(IllegalStateException.class, () -> midjourneyApi.imagine(
  120 + new MidjourneyApi.ImagineRequest(null, "一只猫", null, null)));
  121 + assertEquals("[midjourney-api] 调用失败！", exception.getMessage());
  122 + // 断言：请求
  123 + assertAuthorizedRequest(captureRequest(), "/submit/imagine");
  124 + }
  125 +
  126 + // ========== action ==========
  127 +
  128 + @Test
  129 + public void testAction_success() {
  130 + // mock 方法（Midjourney Proxy 返回）
  131 + mockResponse(HttpStatus.OK, MediaType.APPLICATION_JSON_VALUE,
  132 + "{\"code\":\"1\",\"description\":\"提交成功\",\"result\":\"task-2\"}");
  133 + // 准备参数
  134 + MidjourneyApi.ActionRequest request = new MidjourneyApi.ActionRequest("task-1",
  135 + "MJ::JOB::upsample::1::task-1", null);
  136 +
  137 + // 调用
  138 + MidjourneyApi.SubmitResponse response = midjourneyApi.action(request);
  139 + // 断言：返回
  140 + assertEquals("1", response.code());
  141 + assertEquals("task-2", response.result());
  142 + assertEquals(NOTIFY_URL, request.getNotifyHook());
  143 + // 断言：请求
  144 + assertAuthorizedRequest(captureRequest(), "/submit/action");
  145 + }
  146 +
  147 + @ParameterizedTest
  148 + @ValueSource(ints = {401, 500})
  149 + public void testAction_errorWithoutBody(int status) {
  150 + // mock 方法（Midjourney Proxy 失败，且无响应 body）
  151 + mockResponseWithoutBody(status);
  152 +
  153 + // 调用，并断言异常
  154 + IllegalStateException exception = assertThrows(IllegalStateException.class, () -> midjourneyApi.action(
  155 + new MidjourneyApi.ActionRequest("task-1", "MJ::JOB::upsample::1::task-1", null)));
  156 + assertEquals("[midjourney-api] 调用失败！", exception.getMessage());
  157 + // 断言：请求
  158 + assertAuthorizedRequest(captureRequest(), "/submit/action");
  159 + }
  160 +
  161 + // ========== getTaskList ==========
  162 +
  163 + @Test
  164 + public void testGetTaskList_success() {
  165 + // mock 方法（Midjourney Proxy 返回）
  166 + mockResponse(HttpStatus.OK, MediaType.APPLICATION_JSON_VALUE,
  167 + "[{\"id\":\"task-1\",\"status\":\"SUCCESS\",\"imageUrl\":\"https://cdn.example.com/task-1.png\"}]");
  168 +
  169 + // 调用
  170 + List<MidjourneyApi.Notify> taskList = midjourneyApi.getTaskList(Collections.singletonList("task-1"));
  171 + // 断言：返回
  172 + assertEquals(1, taskList.size());
  173 + assertEquals("task-1", taskList.get(0).id());
  174 + assertEquals("SUCCESS", taskList.get(0).status());
  175 + assertEquals("https://cdn.example.com/task-1.png", taskList.get(0).imageUrl());
  176 + // 断言：请求
  177 + assertAuthorizedRequest(captureRequest(), "/task/list-by-condition");
  178 + }
  179 +
  180 + @Test
  181 + public void testGetTaskList_successWithoutBody() {
  182 + // mock 方法（Midjourney Proxy 成功，但无响应 body）
  183 + mockResponseWithoutBody(200);
  184 +
  185 + // 调用
  186 + List<MidjourneyApi.Notify> taskList = midjourneyApi.getTaskList(Collections.singletonList("task-1"));
  187 + // 断言：返回（保持原有行为，返回空列表）
  188 + assertTrue(taskList.isEmpty());
  189 + // 断言：请求
  190 + assertAuthorizedRequest(captureRequest(), "/task/list-by-condition");
  191 + }
  192 +
  193 + @ParameterizedTest
  194 + @ValueSource(ints = {401, 500})
  195 + public void testGetTaskList_errorWithoutBody(int status) {
  196 + // mock 方法（Midjourney Proxy 失败，且无响应 body）
  197 + mockResponseWithoutBody(status);
  198 +
  199 + // 调用，并断言异常（不能被当成空列表吞掉）
  200 + IllegalStateException exception = assertThrows(IllegalStateException.class,
  201 + () -> midjourneyApi.getTaskList(Collections.singletonList("task-1")));
  202 + assertEquals("[midjourney-api] 调用失败！", exception.getMessage());
  203 + // 断言：请求
  204 + assertAuthorizedRequest(captureRequest(), "/task/list-by-condition");
  205 + }
  206 +
  207 + // ========== 辅助方法 ==========
  208 +
  209 + /**
  210 + * 构造真实的 {@link MidjourneyApi}，仅在构造期间将 {@link WebClient#builder()} 替换为使用 mock {@link ExchangeFunction} 的 builder
  211 + */
  212 + private MidjourneyApi createMidjourneyApi(String apiKey) {
  213 + WebClient.Builder builder = WebClient.builder().exchangeFunction(exchangeFunction);
  214 + try (MockedStatic<WebClient> webClientMock = mockStatic(WebClient.class)) {
  215 + webClientMock.when(WebClient::builder).thenReturn(builder);
  216 + return new MidjourneyApi(BASE_URL, apiKey, NOTIFY_URL);
  217 + }
  218 + }
  219 +
  220 + private void mockResponse(HttpStatus status, String contentType, String body) {
  221 + when(exchangeFunction.exchange(any(ClientRequest.class))).thenReturn(Mono.just(
  222 + ClientResponse.create(status).header(HttpHeaders.CONTENT_TYPE, contentType).body(body).build()));
  223 + }
  224 +
  225 + /**
  226 + * mock 无响应 body 的返回，即 body 为空 Flux，而不是空字符串
  227 + */
  228 + private void mockResponseWithoutBody(int status) {
  229 + when(exchangeFunction.exchange(any(ClientRequest.class))).thenReturn(Mono.just(
  230 + ClientResponse.create(HttpStatus.valueOf(status)).build()));
  231 + }
  232 +
  233 + private ClientRequest captureRequest() {
  234 + ArgumentCaptor<ClientRequest> requestCaptor = ArgumentCaptor.forClass(ClientRequest.class);
  235 + verify(exchangeFunction).exchange(requestCaptor.capture());
  236 + return requestCaptor.getValue();
  237 + }
  238 +
  239 + private static void assertAuthorizedRequest(ClientRequest request, String path) {
  240 + assertEquals(HttpMethod.POST, request.method());
  241 + assertEquals(URI.create("https://mj-proxy.example.com/mj" + path), request.url());
  242 + assertEquals(MediaType.APPLICATION_JSON, request.headers().getContentType());
  243 + assertEquals("Bearer test-mj-api-secret", request.headers().getFirst(HttpHeaders.AUTHORIZATION));
  244 + assertEquals("test-mj-api-secret", request.headers().getFirst("mj-api-secret"));
  245 + }
  246 +
  247 +}
... ...
yudao-module-ai/src/test/java/cn/iocoder/yudao/module/ai/service/image/AiImageServiceImplUnitTest.java 0 → 100644
  View file @2c78172
  1 +package cn.iocoder.yudao.module.ai.service.image;
  2 +
  3 +import cn.iocoder.yudao.framework.test.core.ut.BaseMockitoUnitTest;
  4 +import cn.iocoder.yudao.module.ai.controller.admin.image.vo.midjourney.AiMidjourneyActionReqVO;
  5 +import cn.iocoder.yudao.module.ai.controller.admin.image.vo.midjourney.AiMidjourneyImagineReqVO;
  6 +import cn.iocoder.yudao.module.ai.dal.dataobject.image.AiImageDO;
  7 +import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiModelDO;
  8 +import cn.iocoder.yudao.module.ai.dal.mysql.image.AiImageMapper;
  9 +import cn.iocoder.yudao.module.ai.enums.image.AiImageStatusEnum;
  10 +import cn.iocoder.yudao.module.ai.enums.model.AiPlatformEnum;
  11 +import cn.iocoder.yudao.module.ai.framework.ai.core.model.midjourney.api.MidjourneyApi;
  12 +import cn.iocoder.yudao.module.ai.service.model.AiModelService;
  13 +import cn.iocoder.yudao.module.infra.api.file.FileApi;
  14 +import org.junit.jupiter.api.Test;
  15 +import org.junit.jupiter.params.ParameterizedTest;
  16 +import org.junit.jupiter.params.provider.Arguments;
  17 +import org.junit.jupiter.params.provider.MethodSource;
  18 +import org.mockito.InjectMocks;
  19 +import org.mockito.Mock;
  20 +
  21 +import java.util.Arrays;
  22 +import java.util.Collections;
  23 +import java.util.Map;
  24 +import java.util.Objects;
  25 +import java.util.Set;
  26 +import java.util.stream.Stream;
  27 +
  28 +import static cn.iocoder.yudao.framework.test.core.util.AssertUtils.assertServiceException;
  29 +import static cn.iocoder.yudao.module.ai.enums.ErrorCodeConstants.*;
  30 +import static org.junit.jupiter.api.Assertions.assertEquals;
  31 +import static org.mockito.ArgumentMatchers.any;
  32 +import static org.mockito.ArgumentMatchers.argThat;
  33 +import static org.mockito.Mockito.*;
  34 +
  35 +/**
  36 + * {@link AiImageServiceImpl} 的单元测试类
  37 + *
  38 + * 重点覆盖 Midjourney 任务与模型（Midjourney Proxy）的绑定：
  39 + * 1. {@link AiImageServiceImpl#midjourneyAction(Long, AiMidjourneyActionReqVO)} 生成的 image，继承原 image 的模型
  40 + * 2. {@link AiImageServiceImpl#midjourneySync()} 按照模型分组，使用各自的 API 查询任务进展
  41 + * 3. {@link AiImageServiceImpl#midjourneyImagine(Long, AiMidjourneyImagineReqVO)} 绘图参数（--ar、--v 等）拼接到 prompt 中提交
  42 + *
  43 + * @author 芋道源码
  44 + */
  45 +public class AiImageServiceImplUnitTest extends BaseMockitoUnitTest {
  46 +
  47 + private static final Long USER_ID = 1L;
  48 + private static final Long IMAGE_ID = 10L;
  49 + private static final Long MODEL_ID = 101L;
  50 + private static final Long OTHER_MODEL_ID = 202L;
  51 + private static final String CUSTOM_ID = "MJ::JOB::upsample::1::task-a";
  52 + private static final String NEW_TASK_ID = "task-new";
  53 +
  54 + @InjectMocks
  55 + private AiImageServiceImpl imageService;
  56 +
  57 + @Mock
  58 + private AiModelService modelService;
  59 + @Mock
  60 + private AiImageMapper imageMapper;
  61 + @Mock
  62 + private FileApi fileApi;
  63 +
  64 + @Mock
  65 + private MidjourneyApi midjourneyApi;
  66 + @Mock
  67 + private MidjourneyApi otherMidjourneyApi;
  68 +
  69 + // ========== midjourneyImagine ==========
  70 +
  71 + @ParameterizedTest(name = "{0}")
  72 + @MethodSource("imagineModels")
  73 + public void testMidjourneyImagine_success(String modelCode, String version, String expectedVersionParam,
  74 + String unexpectedParam, String expectedState) {
  75 + // mock 方法（模型）
  76 + AiModelDO model = mockImagineModel(modelCode);
  77 + when(midjourneyApi.imagine(any(MidjourneyApi.ImagineRequest.class))).thenReturn(
  78 + new MidjourneyApi.SubmitResponse(MidjourneyApi.SubmitCodeEnum.SUBMIT_SUCCESS.getCode(), "提交成功", null, NEW_TASK_ID));
  79 +
  80 + // 调用
  81 + imageService.midjourneyImagine(USER_ID, buildImagineReqVO(version));
  82 + // 断言：绘图参数拼接到 prompt 中提交，state 保持原自定义参数
  83 + verify(midjourneyApi).imagine(argThat(request -> request.getPrompt().startsWith("一只猫 ")
  84 + && request.getPrompt().contains("--ar 1024:768")
  85 + && request.getPrompt().contains(expectedVersionParam)
  86 + && !request.getPrompt().contains(unexpectedParam)
  87 + && Objects.equals(request.getState(), expectedState)
  88 + && request.getBase64Array() == null && request.getNotifyHook() == null));
  89 + // 断言：保存的 image 仍是用户原始 prompt，以及模型绑定
  90 + verify(imageMapper).insert(argThat((AiImageDO image) -> Objects.equals(image.getPrompt(), "一只猫")
  91 + && Objects.equals(image.getUserId(), USER_ID)
  92 + && Objects.equals(image.getModelId(), MODEL_ID)
  93 + && Objects.equals(image.getModel(), model.getName())
  94 + && Objects.equals(image.getPlatform(), AiPlatformEnum.MIDJOURNEY.getPlatform())
  95 + && Objects.equals(image.getWidth(), 1024) && Objects.equals(image.getHeight(), 768)
  96 + && Objects.equals(image.getStatus(), AiImageStatusEnum.IN_PROGRESS.getStatus())
  97 + && Boolean.FALSE.equals(image.getPublicStatus())));
  98 + verify(imageMapper).updateById(argThat((AiImageDO update) -> Objects.equals(update.getTaskId(), NEW_TASK_ID)
  99 + && update.getPrompt() == null));
  100 + verifyNoInteractions(fileApi);
  101 + }
  102 +
  103 + private static Stream<Arguments> imagineModels() {
  104 + return Stream.of(
  105 + Arguments.of("midjourney", "6.0", "--v 6.0", "--niji", " --ar 1024:768 --v 6.0 "),
  106 + Arguments.of("niji", "6", "--niji 6", "--v ", " --ar 1024:768 --niji 6 "));
  107 + }
  108 +
  109 + @Test
  110 + public void testMidjourneyImagine_submitFail() {
  111 + // mock 方法：Midjourney Proxy 拒绝
  112 + mockImagineModel("midjourney");
  113 + when(midjourneyApi.imagine(any(MidjourneyApi.ImagineRequest.class))).thenReturn(
  114 + new MidjourneyApi.SubmitResponse("4", "提示词不合法", null, null));
  115 +
  116 + // 调用，并断言异常
  117 + assertServiceException(() -> imageService.midjourneyImagine(USER_ID, buildImagineReqVO("6.0")),
  118 + IMAGE_MIDJOURNEY_SUBMIT_FAIL, "提示词不合法");
  119 + // 断言：只保存了 image，没有更新 taskId
  120 + verify(imageMapper).insert(any(AiImageDO.class));
  121 + verifyNoMoreInteractions(imageMapper);
  122 + verifyNoInteractions(fileApi);
  123 + }
  124 +
  125 + private AiModelDO mockImagineModel(String modelCode) {
  126 + AiModelDO model = new AiModelDO().setId(MODEL_ID).setName("Midjourney 绘图").setModel(modelCode)
  127 + .setPlatform(AiPlatformEnum.MIDJOURNEY.getPlatform());
  128 + when(modelService.validateModel(MODEL_ID)).thenReturn(model);
  129 + when(modelService.getMidjourneyApi(MODEL_ID)).thenReturn(midjourneyApi);
  130 + return model;
  131 + }
  132 +
  133 + private static AiMidjourneyImagineReqVO buildImagineReqVO(String version) {
  134 + return new AiMidjourneyImagineReqVO().setPrompt("一只猫").setModelId(MODEL_ID)
  135 + .setWidth(1024).setHeight(768).setVersion(version); // referImageUrl 为空，不下载参考图
  136 + }
  137 +
  138 + // ========== midjourneyAction ==========
  139 +
  140 + @Test
  141 + public void testMidjourneyAction_success() {
  142 + // mock 数据
  143 + AiImageDO image = buildActionImage(USER_ID);
  144 + when(imageMapper.selectById(IMAGE_ID)).thenReturn(image);
  145 + // mock 方法
  146 + when(modelService.getMidjourneyApi(MODEL_ID)).thenReturn(midjourneyApi);
  147 + when(midjourneyApi.action(any(MidjourneyApi.ActionRequest.class))).thenReturn(
  148 + new MidjourneyApi.SubmitResponse(MidjourneyApi.SubmitCodeEnum.SUBMIT_SUCCESS.getCode(), "提交成功", null, NEW_TASK_ID));
  149 +
  150 + // 调用
  151 + imageService.midjourneyAction(USER_ID, buildActionReqVO());
  152 + // 断言：提交原 image 的任务编号，以及所选按钮的 customId
  153 + verify(midjourneyApi).action(argThat(request -> Objects.equals(request.getTaskId(), image.getTaskId())
  154 + && Objects.equals(request.getCustomId(), CUSTOM_ID)
  155 + && request.getNotifyHook() == null));
  156 + // 断言：新 image 继承原 image 的模型，以及其它绑定字段
  157 + verify(imageMapper).insert(argThat((AiImageDO newImage) -> newImage.getId() == null
  158 + && Objects.equals(newImage.getModelId(), MODEL_ID)
  159 + && Objects.equals(newImage.getModel(), image.getModel())
  160 + && Objects.equals(newImage.getUserId(), USER_ID)
  161 + && Objects.equals(newImage.getPrompt(), image.getPrompt())
  162 + && Objects.equals(newImage.getPlatform(), AiPlatformEnum.MIDJOURNEY.getPlatform())
  163 + && Objects.equals(newImage.getWidth(), image.getWidth())
  164 + && Objects.equals(newImage.getHeight(), image.getHeight())
  165 + && Objects.equals(newImage.getOptions(), image.getOptions())
  166 + && Objects.equals(newImage.getTaskId(), NEW_TASK_ID)
  167 + && Objects.equals(newImage.getStatus(), AiImageStatusEnum.IN_PROGRESS.getStatus())
  168 + && Boolean.FALSE.equals(newImage.getPublicStatus())));
  169 + }
  170 +
  171 + @Test
  172 + public void testMidjourneyAction_notOwner() {
  173 + // mock 数据：其它用户的 image
  174 + when(imageMapper.selectById(IMAGE_ID)).thenReturn(buildActionImage(USER_ID + 1));
  175 +
  176 + // 调用，并断言异常
  177 + assertServiceException(() -> imageService.midjourneyAction(USER_ID, buildActionReqVO()), IMAGE_NOT_EXISTS);
  178 + assertActionNotInserted();
  179 + verifyNoInteractions(modelService);
  180 + }
  181 +
  182 + @Test
  183 + public void testMidjourneyAction_customIdNotExists() {
  184 + // mock 数据
  185 + when(imageMapper.selectById(IMAGE_ID)).thenReturn(buildActionImage(USER_ID));
  186 + // mock 方法
  187 + when(modelService.getMidjourneyApi(MODEL_ID)).thenReturn(midjourneyApi);
  188 +
  189 + // 调用，并断言异常
  190 + assertServiceException(() -> imageService.midjourneyAction(USER_ID, buildActionReqVO().setCustomId("MJ::JOB::other")),
  191 + IMAGE_CUSTOM_ID_NOT_EXISTS);
  192 + assertActionNotInserted();
  193 + verifyNoInteractions(midjourneyApi);
  194 + }
  195 +
  196 + @Test
  197 + public void testMidjourneyAction_submitFail() {
  198 + // mock 数据
  199 + when(imageMapper.selectById(IMAGE_ID)).thenReturn(buildActionImage(USER_ID));
  200 + // mock 方法：Midjourney Proxy 拒绝
  201 + when(modelService.getMidjourneyApi(MODEL_ID)).thenReturn(midjourneyApi);
  202 + when(midjourneyApi.action(any(MidjourneyApi.ActionRequest.class))).thenReturn(
  203 + new MidjourneyApi.SubmitResponse("4", "任务不存在", null, null));
  204 +
  205 + // 调用，并断言异常
  206 + assertServiceException(() -> imageService.midjourneyAction(USER_ID, buildActionReqVO()),
  207 + IMAGE_MIDJOURNEY_SUBMIT_FAIL, "任务不存在");
  208 + assertActionNotInserted();
  209 + }
  210 +
  211 + private AiImageDO buildActionImage(Long userId) {
  212 + return new AiImageDO().setId(IMAGE_ID).setUserId(userId).setPrompt("一只猫")
  213 + .setPlatform(AiPlatformEnum.MIDJOURNEY.getPlatform()).setModelId(MODEL_ID).setModel("midjourney")
  214 + .setWidth(1024).setHeight(768).setOptions(Map.of("version", "6.0"))
  215 + .setStatus(AiImageStatusEnum.SUCCESS.getStatus()).setPublicStatus(true).setTaskId("task-a")
  216 + .setButtons(Collections.singletonList(new MidjourneyApi.Button(CUSTOM_ID, null, "U1", null, null)));
  217 + }
  218 +
  219 + private static AiMidjourneyActionReqVO buildActionReqVO() {
  220 + return new AiMidjourneyActionReqVO().setId(IMAGE_ID).setCustomId(CUSTOM_ID);
  221 + }
  222 +
  223 + /**
  224 + * 断言：只查询了 image，没有新增 image
  225 + */
  226 + private void assertActionNotInserted() {
  227 + verify(imageMapper).selectById(IMAGE_ID);
  228 + verifyNoMoreInteractions(imageMapper);
  229 + }
  230 +
  231 + // ========== midjourneySync ==========
  232 +
  233 + @Test
  234 + public void testMidjourneySync_empty() {
  235 + // mock 数据
  236 + mockInProgressImages();
  237 +
  238 + // 调用，并断言
  239 + assertEquals(0, imageService.midjourneySync());
  240 + verifyNoInteractions(modelService);
  241 + }
  242 +
  243 + @Test
  244 + public void testMidjourneySync_singleModel() {
  245 + // mock 数据
  246 + mockInProgressImages(buildSyncImage(1L, MODEL_ID, "task-a"), buildSyncImage(2L, MODEL_ID, "task-b"));
  247 + // mock 方法
  248 + when(modelService.getMidjourneyApi(MODEL_ID)).thenReturn(midjourneyApi);
  249 + when(midjourneyApi.getTaskList(Set.of("task-a", "task-b"))).thenReturn(Arrays.asList(
  250 + buildNotify("task-a", MidjourneyApi.TaskStatusEnum.SUCCESS, null),
  251 + buildNotify("task-b", MidjourneyApi.TaskStatusEnum.FAILURE, "被拒绝")));
  252 +
  253 + // 调用，并断言
  254 + assertEquals(2, imageService.midjourneySync());
  255 + verifyImageUpdated(1L, AiImageStatusEnum.SUCCESS, null);
  256 + verifyImageUpdated(2L, AiImageStatusEnum.FAIL, "被拒绝");
  257 + assertSyncNoMoreWrite();
  258 + }
  259 +
  260 + @Test
  261 + public void testMidjourneySync_multiModel() {
  262 + // mock 数据：两个模型，各自对应不同的 Midjourney Proxy
  263 + mockInProgressImages(buildSyncImage(1L, MODEL_ID, "task-a"), buildSyncImage(2L, OTHER_MODEL_ID, "task-b"));
  264 + // mock 方法：每个 API 只查询自己模型的任务
  265 + when(modelService.getMidjourneyApi(MODEL_ID)).thenReturn(midjourneyApi);
  266 + when(midjourneyApi.getTaskList(Set.of("task-a"))).thenReturn(Collections.singletonList(
  267 + buildNotify("task-a", MidjourneyApi.TaskStatusEnum.SUCCESS, null)));
  268 + when(modelService.getMidjourneyApi(OTHER_MODEL_ID)).thenReturn(otherMidjourneyApi);
  269 + when(otherMidjourneyApi.getTaskList(Set.of("task-b"))).thenReturn(Collections.singletonList(
  270 + buildNotify("task-b", MidjourneyApi.TaskStatusEnum.SUCCESS, null)));
  271 +
  272 + // 调用，并断言
  273 + assertEquals(2, imageService.midjourneySync());
  274 + verifyImageUpdated(1L, AiImageStatusEnum.SUCCESS, null);
  275 + verifyImageUpdated(2L, AiImageStatusEnum.SUCCESS, null);
  276 + assertSyncNoMoreWrite();
  277 + }
  278 +
  279 + @Test
  280 + public void testMidjourneySync_sameTaskIdInDifferentModels() {
  281 + // mock 数据：不同 Midjourney Proxy 中，相同的 taskId
  282 + mockInProgressImages(buildSyncImage(1L, MODEL_ID, "task-same"), buildSyncImage(2L, OTHER_MODEL_ID, "task-same"));
  283 + // mock 方法：两个 Proxy 返回不同的结果
  284 + when(modelService.getMidjourneyApi(MODEL_ID)).thenReturn(midjourneyApi);
  285 + when(midjourneyApi.getTaskList(Set.of("task-same"))).thenReturn(Collections.singletonList(
  286 + buildNotify("task-same", MidjourneyApi.TaskStatusEnum.SUCCESS, null)));
  287 + when(modelService.getMidjourneyApi(OTHER_MODEL_ID)).thenReturn(otherMidjourneyApi);
  288 + when(otherMidjourneyApi.getTaskList(Set.of("task-same"))).thenReturn(Collections.singletonList(
  289 + buildNotify("task-same", MidjourneyApi.TaskStatusEnum.FAILURE, "被拒绝")));
  290 +
  291 + // 调用，并断言：各自使用自己 Proxy 的结果，不串
  292 + assertEquals(2, imageService.midjourneySync());
  293 + verifyImageUpdated(1L, AiImageStatusEnum.SUCCESS, null);
  294 + verifyImageUpdated(2L, AiImageStatusEnum.FAIL, "被拒绝");
  295 + assertSyncNoMoreWrite();
  296 + }
  297 +
  298 + @Test
  299 + public void testMidjourneySync_unboundImageFirst() {
  300 + // mock 数据：未关联模型的 image 排在第一个
  301 + mockInProgressImages(buildSyncImage(1L, null, "task-x"), buildSyncImage(2L, MODEL_ID, "task-a"));
  302 + // mock 方法
  303 + when(modelService.getMidjourneyApi(MODEL_ID)).thenReturn(midjourneyApi);
  304 + when(midjourneyApi.getTaskList(Set.of("task-a"))).thenReturn(Collections.singletonList(
  305 + buildNotify("task-a", MidjourneyApi.TaskStatusEnum.SUCCESS, null)));
  306 +
  307 + // 调用，并断言：跳过未关联模型的 image，其它 image 继续同步
  308 + assertEquals(1, imageService.midjourneySync());
  309 + verifyImageUpdated(2L, AiImageStatusEnum.SUCCESS, null);
  310 + assertSyncNoMoreWrite();
  311 + verifyNoMoreInteractions(modelService);
  312 + }
  313 +
  314 + @Test
  315 + public void testMidjourneySync_allUnbound() {
  316 + // mock 数据：全部未关联模型
  317 + mockInProgressImages(buildSyncImage(1L, null, "task-x"), buildSyncImage(2L, null, "task-y"));
  318 +
  319 + // 调用，并断言：不查询 API，不更新 image
  320 + assertEquals(0, imageService.midjourneySync());
  321 + verifyNoInteractions(modelService);
  322 + assertSyncNoMoreWrite();
  323 + }
  324 +
  325 + @Test
  326 + public void testMidjourneySync_taskNotFound() {
  327 + // mock 数据
  328 + mockInProgressImages(buildSyncImage(1L, MODEL_ID, "task-a"), buildSyncImage(2L, MODEL_ID, "task-b"));
  329 + // mock 方法：Proxy 只返回 task-a 的进展
  330 + when(modelService.getMidjourneyApi(MODEL_ID)).thenReturn(midjourneyApi);
  331 + when(midjourneyApi.getTaskList(Set.of("task-a", "task-b"))).thenReturn(Collections.singletonList(
  332 + buildNotify("task-a", MidjourneyApi.TaskStatusEnum.SUCCESS, null)));
  333 +
  334 + // 调用，并断言：只跳过查询不到的 image
  335 + assertEquals(1, imageService.midjourneySync());
  336 + verifyImageUpdated(1L, AiImageStatusEnum.SUCCESS, null);
  337 + assertSyncNoMoreWrite();
  338 + }
  339 +
  340 + private void mockInProgressImages(AiImageDO... images) {
  341 + when(imageMapper.selectListByStatusAndPlatform(AiImageStatusEnum.IN_PROGRESS.getStatus(),
  342 + AiPlatformEnum.MIDJOURNEY.getPlatform())).thenReturn(Arrays.asList(images));
  343 + }
  344 +
  345 + private static AiImageDO buildSyncImage(Long id, Long modelId, String taskId) {
  346 + return new AiImageDO().setId(id).setModelId(modelId).setTaskId(taskId)
  347 + .setPlatform(AiPlatformEnum.MIDJOURNEY.getPlatform()).setStatus(AiImageStatusEnum.IN_PROGRESS.getStatus());
  348 + }
  349 +
  350 + /**
  351 + * 构建任务进展。imageUrl 为空，不触发图片下载
  352 + */
  353 + private static MidjourneyApi.Notify buildNotify(String taskId, MidjourneyApi.TaskStatusEnum status, String failReason) {
  354 + return new MidjourneyApi.Notify(taskId, null, status.name(), null, null, null, null,
  355 + null, null, null, null, null, failReason, null);
  356 + }
  357 +
  358 + private void verifyImageUpdated(Long id, AiImageStatusEnum status, String errorMessage) {
  359 + verify(imageMapper).updateById(argThat((AiImageDO update) -> Objects.equals(update.getId(), id)
  360 + && Objects.equals(update.getStatus(), status.getStatus())
  361 + && Objects.equals(update.getErrorMessage(), errorMessage)
  362 + && update.getFinishTime() != null));
  363 + }
  364 +
  365 + /**
  366 + * 断言：除已校验的更新外，没有其它 image 写入，也没有图片下载上传
  367 + */
  368 + private void assertSyncNoMoreWrite() {
  369 + verify(imageMapper).selectListByStatusAndPlatform(AiImageStatusEnum.IN_PROGRESS.getStatus(),
  370 + AiPlatformEnum.MIDJOURNEY.getPlatform());
  371 + verifyNoMoreInteractions(imageMapper);
  372 + verifyNoInteractions(fileApi);
  373 + }
  374 +
  375 +}
... ...