xieyinlu/ruoyi-ai
1
0
Fork 0
mirror of https://gitcode.com/ageerle/ruoyi-ai.git synced 2026-11-11 21:40:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: support Atlas multimodal generation and media preview

Browse Source
This commit is contained in:
ageerle
2026-09-11 10:49:29 +08:00
parent 4766d2e628
commit cca3090577
11 changed files with 498 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
ruoyi-common/ruoyi-common-chat/src/main/java/org/ruoyi/common/chat/security/ChatModelCredentialPolicy.java
View File

@@ -16,6 +16,12 @@ public final class ChatModelCredentialPolicy {
public static final String DEEPSEEK_API_HOST = "https://api.deepseek.com";
public static final String PPIO_PROVIDER = "ppio";
public static final String PPIO_API_HOST = "https://api.ppio.com/openai/v1";
public static final String ATLAS_PROVIDER = "atlas";
private static final Set<String> ATLAS_API_HOSTS = Set.of(
"https://api.atlascloud.ai", "https://api.atlascloud.ai/",
"https://api.atlascloud.ai/v1", "https://api.atlascloud.ai/v1/",
"https://api.atlascloud.ai/api/v1", "https://api.atlascloud.ai/api/v1/"
);
private static final Set<String> PPIO_API_HOSTS = Set.of(
PPIO_API_HOST, PPIO_API_HOST + "/",
@@ -69,6 +75,10 @@ public final class ChatModelCredentialPolicy {
String apiHost,
String apiKey) {
requireSecureApiHost(apiHost);
if (ATLAS_PROVIDER.equals(providerCode)) {
requireAtlasConfiguration(providerCode, modelName, apiHost, apiKey);
return;
}
if (CustomApiCredentialPolicy.isCustomProvider(providerCode)) {
CustomApiCredentialPolicy.requireConfiguration(providerCode, modelName, apiHost, apiKey);
return;
@@ -121,6 +131,8 @@ public final class ChatModelCredentialPolicy {
expectedReference = "env:DEEPSEEK_API_KEY";
} else if (PPIO_PROVIDER.equals(providerCode)) {
expectedReference = "env:PPIO_API_KEY";
} else if (ATLAS_PROVIDER.equals(providerCode)) {
expectedReference = "env:ATLAS_API_KEY";
} else {
throw new IllegalArgumentException("Model credential provider is not allowlisted");
}
@@ -137,11 +149,29 @@ public final class ChatModelCredentialPolicy {
CustomApiCredentialPolicy.requireConfiguration(providerCode, modelName, apiHost, apiKey);
} else if (PPIO_PROVIDER.equals(providerCode)) {
requirePpioConfiguration(providerCode, modelName, apiHost, apiKey);
} else if (ATLAS_PROVIDER.equals(providerCode)) {
requireAtlasConfiguration(providerCode, modelName, apiHost, apiKey);
} else {
requireDeepSeekConfiguration(providerCode, modelName, apiHost, apiKey);
}
}
/** Atlas credentials may only be sent to its official media API origin. */
public static void requireAtlasConfiguration(String providerCode, String modelName,
String apiHost, String apiKey) {
if (!ATLAS_PROVIDER.equals(providerCode)) {
throw new IllegalArgumentException("Atlas model provider is not trusted");
}
if (modelName == null || !modelName.matches("[A-Za-z0-9][A-Za-z0-9._/-]{0,254}")) {
throw new IllegalArgumentException("Atlas model ID is invalid");
}
requireSecureApiHost(apiHost);
if (!ATLAS_API_HOSTS.contains(apiHost)) {
throw new IllegalArgumentException("Atlas API host is not allowlisted");
}
requireProviderReference(providerCode, apiKey);
}
/** Validates PPIO and returns the base URL expected by the OpenAI Chat Completions client. */
public static String requirePpioConfiguration(String providerCode, String modelName,
String apiHost, String apiKey) {

2
ruoyi-common/ruoyi-common-chat/src/main/java/org/ruoyi/common/chat/security/ChatModelSecretReference.java
View File

@@ -8,7 +8,7 @@ import java.util.regex.Pattern;
public final class ChatModelSecretReference {
public static final String ENV_REFERENCE_REGEXP =
"^env:(?:(?:DEEPSEEK|PPIO)_API_KEY|" + CustomApiCredentialPolicy.REFERENCE_PATTERN + ")$";
"^env:(?:(?:DEEPSEEK|PPIO|ATLAS)_API_KEY|" + CustomApiCredentialPolicy.REFERENCE_PATTERN + ")$";
private static final String ENV_PREFIX = "env:";
private static final Pattern ENV_REFERENCE = Pattern.compile(ENV_REFERENCE_REGEXP);

50
ruoyi-common/ruoyi-common-chat/src/test/java/org/ruoyi/common/chat/security/AtlasCredentialPolicyTest.java Normal file
View File

@@ -0,0 +1,50 @@
package org.ruoyi.common.chat.security;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import static org.junit.jupiter.api.Assertions.*;
@Tag("dev")
class AtlasCredentialPolicyTest {
private static final String MODEL = "bytedance/seed-audio-1.0";
private static final String HOST = "https://api.atlascloud.ai/v1";
private static final String REF = "env:ATLAS_API_KEY";
@Test
void acceptsOfficialEndpointFormsAtSaveAndUse() {
for (String suffix : new String[]{"", "/", "/v1", "/v1/", "/api/v1", "/api/v1/"}) {
String host = "https://api.atlascloud.ai" + suffix;
assertDoesNotThrow(() -> ChatModelCredentialPolicy.requirePersistableConfiguration("atlas", MODEL, host, REF));
assertEquals("test-secret", ChatModelCredentialPolicy.resolveApiKeyForUse("atlas", "atlas", MODEL, host, REF,
name -> "ATLAS_API_KEY".equals(name) ? "test-secret" : null));
}
}
@Test
void rejectsUntrustedEndpointsBeforeReadingEnvironment() {
for (String host : new String[]{"http://api.atlascloud.ai/v1", "https://evil.example/v1",
"https://api.atlascloud.ai.evil.example/v1", "https://api.atlascloud.ai:8443/v1",
"https://user@api.atlascloud.ai/v1", "https://api.atlascloud.ai/v1?x=1", "https://api.atlascloud.ai/other"}) {
assertThrows(IllegalArgumentException.class,
() -> ChatModelCredentialPolicy.requirePersistableConfiguration("atlas", MODEL, host, REF));
assertThrows(IllegalArgumentException.class,
() -> ChatModelCredentialPolicy.resolveApiKeyForUse("atlas", "atlas", MODEL, host, REF,
name -> { fail("Must validate endpoint before resolving secret"); return null; }));
}
}
@Test
void rejectsWrongProviderReferenceAndMissingSecret() {
assertThrows(IllegalArgumentException.class, () -> ChatModelCredentialPolicy.resolveApiKeyForUse(
"openai", "atlas", MODEL, HOST, REF, name -> "secret"));
for (String key : new String[]{"env:PPIO_API_KEY", "plaintext", "", null}) {
assertThrows(IllegalArgumentException.class,
() -> ChatModelCredentialPolicy.requirePersistableConfiguration("atlas", MODEL, HOST, key));
}
assertThrows(IllegalStateException.class, () -> ChatModelCredentialPolicy.resolveApiKeyForUse(
"atlas", "atlas", MODEL, HOST, REF, name -> null));
assertThrows(IllegalArgumentException.class,
() -> ChatModelCredentialPolicy.requirePersistableConfiguration("atlas", "bad model", HOST, REF));
}
}

44
ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/controller/chat/MediaContentController.java Normal file
View File

@@ -0,0 +1,44 @@
package org.ruoyi.controller.chat;
import jakarta.validation.constraints.NotBlank;
import lombok.RequiredArgsConstructor;
import org.ruoyi.common.chat.service.chat.IChatModelService;
import org.ruoyi.service.media.AtlasMediaContentService;
import org.springframework.http.CacheControl;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import java.io.IOException;
@Validated
@RestController
@RequiredArgsConstructor
@RequestMapping("/media")
public class MediaContentController {
private final IChatModelService models;
private final AtlasMediaContentService content;
/** Uses the normal login interceptor; browser clients fetch with their Authorization header. */
@GetMapping("/content")
public ResponseEntity<byte[]> content(@RequestParam @NotBlank String model,
@RequestParam @NotBlank String predictionId) throws IOException {
var configuredModel = models.selectModelByName(model);
if (configuredModel == null) {
throw new IllegalArgumentException("未找到模型配置");
}
var result = content.retrieve(configuredModel, predictionId);
return ResponseEntity.ok()
.contentType(MediaType.parseMediaType(result.mimeType()))
.contentLength(result.bytes().length)
.cacheControl(CacheControl.noStore())
.header(HttpHeaders.CONTENT_DISPOSITION, "inline")
.header("X-Content-Type-Options", "nosniff")
.body(result.bytes());
}
}

20
ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/controller/chat/MediaGenerationController.java
View File

@@ -56,14 +56,18 @@ public class MediaGenerationController {
@PostMapping("/image")
public R<MediaGenerationResponse> image(@Valid @RequestBody ImageGenerationRequest request) {
ChatModelVo model = loadModel(request.getModel(), ModelType.IMAGE.getKey());
String result = imageServiceFactory.getOriginalService(model.getProviderCode())
.generateImage(ImageContext.builder()
.chatModelVo(model)
.prompt(request.getPrompt())
.size(request.getSize())
.seed(request.getSeed())
.build());
return R.ok(toImageResponse(result));
ImageContext context = ImageContext.builder()
.chatModelVo(model)
.prompt(request.getPrompt())
.size(request.getSize())
.seed(request.getSeed())
.build();
var service = imageServiceFactory.getOriginalService(model.getProviderCode());
// Atlas sync requests can outlive its gateway timeout. Return a task for the workbench to poll.
if ("atlas".equals(model.getProviderCode())) {
return R.ok(service.startImageGeneration(context));
}
return R.ok(toImageResponse(service.generateImage(context)));
}
@PostMapping("/video")

97
ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/service/media/AtlasMediaContentService.java Normal file
View File

@@ -0,0 +1,97 @@
package org.ruoyi.service.media;
import okhttp3.HttpUrl;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.Response;
import okhttp3.ResponseBody;
import org.ruoyi.common.chat.domain.vo.chat.ChatModelVo;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import java.io.IOException;
import java.util.Locale;
import java.util.Set;
import java.util.concurrent.TimeUnit;
/** Delivers Atlas task outputs without exposing a general-purpose URL proxy. */
@Service
public class AtlasMediaContentService {
static final int MAX_BYTES = 64 * 1024 * 1024;
private static final Set<String> OUTPUT_HOSTS = Set.of(
"atlas-media.oss-us-west-1.aliyuncs.com",
"ark-acg-ap-southeast-1.tos-ap-southeast-1.volces.com"
);
private static final Set<String> MIME_TYPES = Set.of(
"image/jpeg", "image/png", "image/webp", "image/gif",
"audio/mpeg", "audio/wav", "audio/x-wav", "audio/ogg", "audio/mp4", "audio/flac",
"video/mp4", "video/webm"
);
private final AtlasPredictionService predictions;
private final OkHttpClient client;
@Autowired
public AtlasMediaContentService(AtlasPredictionService predictions) {
this(predictions, new OkHttpClient.Builder()
.followRedirects(false).followSslRedirects(false)
.connectTimeout(15, TimeUnit.SECONDS)
.readTimeout(60, TimeUnit.SECONDS)
.callTimeout(90, TimeUnit.SECONDS)
.build());
}
AtlasMediaContentService(AtlasPredictionService predictions, OkHttpClient client) {
this.predictions = predictions;
this.client = client;
}
public Content retrieve(ChatModelVo model, String predictionId) throws IOException {
if (!"atlas".equals(model.getProviderCode())
|| !Set.of("image", "audio", "video").contains(model.getCategory())) {
throw new IllegalArgumentException("此模型不支持 Atlas 资源预览");
}
if (predictionId == null || !predictionId.matches("[A-Za-z0-9_-]{1,128}")) {
throw new IllegalArgumentException("无效的任务 ID");
}
// Resolve the output using this configured model's credentials, never a caller-supplied URL.
var result = predictions.retrieve(model, predictionId);
if (!Set.of("completed", "succeeded", "success").contains(String.valueOf(result.getStatus()))
|| !predictionId.equals(result.getId())) {
throw new IllegalArgumentException("任务尚未完成、已过期或不存在");
}
HttpUrl url = checkedOutputUrl(result.getUrl());
// The API key belongs only on the prediction API request, never on the media CDN request.
Request request = new Request.Builder().url(url).get().build();
try (Response response = client.newCall(request).execute()) {
ResponseBody body = response.body();
if (response.code() != 200 || body == null) {
throw new IOException("资源服务暂不可用，请稍后重试");
}
String mime = response.header("Content-Type", "").split(";", 2)[0].trim().toLowerCase(Locale.ROOT);
if (!MIME_TYPES.contains(mime) || !mime.startsWith(model.getCategory() + "/")) {
throw new IOException("资源格式与模型类型不匹配或不支持预览");
}
if (body.contentLength() > MAX_BYTES) {
throw new IOException("资源超过 64 MiB 预览上限");
}
byte[] bytes = body.byteStream().readNBytes(MAX_BYTES + 1);
if (bytes.length == 0 || bytes.length > MAX_BYTES) {
throw new IOException("资源为空或超过 64 MiB 预览上限");
}
return new Content(bytes, mime);
}
}
static HttpUrl checkedOutputUrl(String value) {
HttpUrl url = value == null ? null : HttpUrl.parse(value);
if (url == null || !url.isHttps() || url.port() != 443
|| !url.username().isEmpty() || !url.password().isEmpty() || url.fragment() != null
|| !OUTPUT_HOSTS.contains(url.host())) {
throw new IllegalArgumentException("资源地址不在 Atlas 预览白名单内");
}
return url;
}
public record Content(byte[] bytes, String mimeType) { }
}

29
ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/service/media/AtlasPredictionService.java
View File

@@ -43,7 +43,8 @@ public class AtlasPredictionService {
log.warn("Atlas Cloud 预测任务不存在或已过期: predictionId={}, category={}", predictionId, model.getCategory());
return MediaGenerationResponse.builder()
.type(mediaType(model.getCategory()))
.mimeType("image".equals(model.getCategory()) ? "image/png" : "video/mp4")
.mimeType(mimeType(mediaType(model.getCategory()), null))
.id(predictionId)
.status("failed")
.rawResponse(responseText)
.build();
@@ -71,7 +72,7 @@ public class AtlasPredictionService {
}
return MediaGenerationResponse.builder()
.type(type)
.mimeType("image".equals(type) ? "image/png" : "video/mp4")
.mimeType(mimeType(type, firstOutput(data)))
.id(AtlasMediaSupport.text(data, "id"))
.status(status)
.url(firstOutput(data))
@@ -80,6 +81,22 @@ public class AtlasPredictionService {
.build();
}
private String mimeType(String type, String url) {
String path = url == null ? "" : url.split("[?#]", 2)[0].toLowerCase(java.util.Locale.ROOT);
if ("audio".equals(type)) {
if (path.endsWith(".wav")) return "audio/wav";
if (path.endsWith(".ogg") || path.endsWith(".opus")) return "audio/ogg";
if (path.endsWith(".pcm")) return "application/octet-stream";
return "audio/mpeg";
}
if ("image".equals(type)) {
if (path.endsWith(".jpg") || path.endsWith(".jpeg")) return "image/jpeg";
if (path.endsWith(".webp")) return "image/webp";
return "image/png";
}
return "video/mp4";
}
/**
* 提取末帧 URL。Atlas return_last_frame=true 时会在 outputs 或顶层节点返回末帧图片，
* 字段名兼容 last_frame_url / end_frame_url / last_frame / last_frame_image。
@@ -128,7 +145,11 @@ public class AtlasPredictionService {
return null;
}
private String mediaType(String category) {
return "image".equals(category) ? "image" : "video";
static String mediaType(String category) {
return switch (category) {
case "image" -> "image";
case "audio" -> "audio";
default -> "video";
};
}
}

29
ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/controller/chat/MediaContentControllerTest.java Normal file
View File

@@ -0,0 +1,29 @@
package org.ruoyi.controller.chat;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.ruoyi.common.chat.domain.vo.chat.ChatModelVo;
import org.ruoyi.common.chat.service.chat.IChatModelService;
import org.ruoyi.service.media.AtlasMediaContentService;
import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.Mockito.*;
@Tag("dev")
class MediaContentControllerTest {
@Test
void returnsInlineNonCachedBytes() throws Exception {
var models = mock(IChatModelService.class);
var service = mock(AtlasMediaContentService.class);
var model = new ChatModelVo();
when(models.selectModelByName("image-model")).thenReturn(model);
when(service.retrieve(model, "task-123")).thenReturn(new AtlasMediaContentService.Content(new byte[]{1, 2}, "image/jpeg"));
var result = new MediaContentController(models, service).content("image-model", "task-123");
assertEquals("image/jpeg", result.getHeaders().getContentType().toString());
assertEquals("inline", result.getHeaders().getFirst("Content-Disposition"));
assertEquals("no-store", result.getHeaders().getCacheControl());
assertEquals("nosniff", result.getHeaders().getFirst("X-Content-Type-Options"));
assertEquals(2, result.getHeaders().getContentLength());
assertArrayEquals(new byte[]{1, 2}, result.getBody());
}
}

65
ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/controller/chat/MediaGenerationControllerTest.java Normal file
View File

@@ -0,0 +1,65 @@
package org.ruoyi.controller.chat;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.ruoyi.common.chat.domain.vo.chat.ChatModelVo;
import org.ruoyi.common.chat.entity.media.MediaGenerationResponse;
import org.ruoyi.common.chat.factory.ImageServiceFactory;
import org.ruoyi.common.chat.service.chat.IChatModelService;
import org.ruoyi.common.chat.service.image.IImageGenerationService;
import org.ruoyi.domain.bo.media.ImageGenerationRequest;
import org.ruoyi.service.media.AtlasPredictionService;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;
@Tag("dev")
class MediaGenerationControllerTest {
@Test
void returnsAtlasTaskWithoutBlockingForImage() {
var models = mock(IChatModelService.class);
var factory = mock(ImageServiceFactory.class);
var service = mock(IImageGenerationService.class);
var model = model("atlas");
when(models.selectModelByName("test-image")).thenReturn(model);
when(factory.getOriginalService("atlas")).thenReturn(service);
when(service.startImageGeneration(any())).thenReturn(MediaGenerationResponse.builder()
.id("task-123").type("image").status("pending").build());
var controller = new MediaGenerationController(models, null, factory, null, new AtlasPredictionService());
var result = controller.image(request()).getData();
assertEquals("task-123", result.getId());
assertEquals("pending", result.getStatus());
verify(service, never()).generateImage(any());
}
@Test
void preservesOtherProvidersBase64ImageResponse() {
var models = mock(IChatModelService.class);
var factory = mock(ImageServiceFactory.class);
var service = mock(IImageGenerationService.class);
when(models.selectModelByName("test-image")).thenReturn(model("openai"));
when(factory.getOriginalService("openai")).thenReturn(service);
when(service.generateImage(any())).thenReturn("data:image/png;base64,aGVsbG8=");
var controller = new MediaGenerationController(models, null, factory, null, new AtlasPredictionService());
var result = controller.image(request()).getData();
assertEquals("image/png", result.getMimeType());
assertEquals("aGVsbG8=", result.getB64Json());
verify(service, never()).startImageGeneration(any());
}
private ChatModelVo model(String provider) {
var model = new ChatModelVo();
model.setModelName("test-image");
model.setProviderCode(provider);
model.setCategory("image");
return model;
}
private ImageGenerationRequest request() {
var request = new ImageGenerationRequest();
request.setModel("test-image");
request.setPrompt("a blue book");
return request;
}
}

119
ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/service/media/AtlasMediaContentServiceTest.java Normal file
View File

@@ -0,0 +1,119 @@
package org.ruoyi.service.media;
import okhttp3.Call;
import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.Protocol;
import okhttp3.Request;
import okhttp3.Response;
import okhttp3.ResponseBody;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;
import org.ruoyi.common.chat.domain.vo.chat.ChatModelVo;
import org.ruoyi.common.chat.entity.media.MediaGenerationResponse;
import java.io.IOException;
import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;
@Tag("dev")
class AtlasMediaContentServiceTest {
private static final String URL = "https://atlas-media.oss-us-west-1.aliyuncs.com/assetd-history/result.mp3";
private final AtlasPredictionService predictions = mock(AtlasPredictionService.class);
private final OkHttpClient client = mock(OkHttpClient.class);
private final AtlasMediaContentService service = new AtlasMediaContentService(predictions, client);
@Test
void servesMediaBytesWithoutForwardingCredentialsOrDownloadHeaders() throws Exception {
var model = model();
model.setApiKey("env:ATLAS_API_KEY");
completed(model);
respond(200, ResponseBody.create("ID3audio", MediaType.get("audio/mpeg")), "audio/mpeg");
var content = service.retrieve(model, "task-123");
assertEquals("audio/mpeg", content.mimeType());
assertArrayEquals("ID3audio".getBytes(), content.bytes());
var capture = ArgumentCaptor.forClass(Request.class);
verify(client).newCall(capture.capture());
assertEquals(URL, capture.getValue().url().toString());
assertNull(capture.getValue().header("Authorization"));
assertNull(capture.getValue().header("Cookie"));
}
@Test
void rejectsUntrustedUrlsAndNonHttpsVariants() {
for (String url : new String[]{"http://atlas-media.oss-us-west-1.aliyuncs.com/a.mp3",
"https://127.0.0.1/a", "https://atlas-media.oss-us-west-1.aliyuncs.com.evil.test/a",
"https://user@atlas-media.oss-us-west-1.aliyuncs.com/a",
"https://atlas-media.oss-us-west-1.aliyuncs.com:8443/a",
"https://atlas-media.oss-us-west-1.aliyuncs.com/a#fragment", "file:///etc/passwd"}) {
assertThrows(IllegalArgumentException.class, () -> AtlasMediaContentService.checkedOutputUrl(url));
}
assertEquals("atlas-media.oss-us-west-1.aliyuncs.com", AtlasMediaContentService.checkedOutputUrl(URL).host());
}
@Test
void rejectsInvalidIdsAndOtherProvidersBeforeMakingRequests() {
var model = model();
assertThrows(IllegalArgumentException.class, () -> service.retrieve(model, "../task?url=bad"));
model.setProviderCode("openai");
assertThrows(IllegalArgumentException.class, () -> service.retrieve(model, "task-123"));
verifyNoInteractions(predictions, client);
}
@Test
void rejectsUnfinishedOrMismatchedTasksBeforeFetchingContent() {
var model = model();
when(predictions.retrieve(model, "task-123")).thenReturn(MediaGenerationResponse.builder()
.id("task-123").status("processing").url(URL).build());
assertThrows(IllegalArgumentException.class, () -> service.retrieve(model, "task-123"));
when(predictions.retrieve(model, "task-123")).thenReturn(MediaGenerationResponse.builder()
.id("different").status("completed").url(URL).build());
assertThrows(IllegalArgumentException.class, () -> service.retrieve(model, "task-123"));
verifyNoInteractions(client);
}
@Test
void rejectsRedirectsAndHtmlResponses() throws Exception {
var model = model();
completed(model);
respond(302, ResponseBody.create("", MediaType.get("audio/mpeg")), "audio/mpeg");
assertThrows(IOException.class, () -> service.retrieve(model, "task-123"));
respond(200, ResponseBody.create("<html>error</html>", MediaType.get("text/html")), "text/html");
assertThrows(IOException.class, () -> service.retrieve(model, "task-123"));
}
@Test
void rejectsOversizedContentBeforeReadingTheBody() throws Exception {
var model = model();
completed(model);
var body = mock(ResponseBody.class);
when(body.contentLength()).thenReturn((long) AtlasMediaContentService.MAX_BYTES + 1);
respond(200, body, "audio/mpeg");
assertThrows(IOException.class, () -> service.retrieve(model, "task-123"));
verify(body, never()).byteStream();
}
private ChatModelVo model() {
var model = new ChatModelVo();
model.setCategory("audio");
model.setProviderCode("atlas");
return model;
}
private void completed(ChatModelVo model) {
when(predictions.retrieve(model, "task-123")).thenReturn(MediaGenerationResponse.builder()
.id("task-123").status("completed").url(URL).build());
}
private void respond(int status, ResponseBody body, String type) throws IOException {
Call call = mock(Call.class);
when(client.newCall(any())).thenReturn(call);
when(call.execute()).thenReturn(new Response.Builder().request(new Request.Builder().url(URL).build())
.protocol(Protocol.HTTP_1_1).code(status).message("test").body(body)
.header("Content-Type", type).header("Content-Disposition", "attachment")
.header("x-oss-force-download", "true").header("Location", "http://127.0.0.1/private").build());
}
}

26
ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/service/media/AtlasPredictionServiceTest.java Normal file
View File

@@ -0,0 +1,26 @@
package org.ruoyi.service.media;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import static org.junit.jupiter.api.Assertions.assertEquals;
@Tag("dev")
class AtlasPredictionServiceTest {
@Test
void keepsAudioTypeAndMimeThroughPolling() throws Exception {
var service = new AtlasPredictionService();
assertEquals("audio", AtlasPredictionService.mediaType("audio"));
assertEquals("image", AtlasPredictionService.mediaType("image"));
assertEquals("video", AtlasPredictionService.mediaType("video"));
var pending = service.toResponse("{\"data\":{\"id\":\"task-123\",\"status\":\"processing\"}}", "audio");
assertEquals("audio", pending.getType());
assertEquals("audio/mpeg", pending.getMimeType());
assertEquals("task-123", pending.getId());
for (String[] sample : new String[][]{{"mp3", "audio/mpeg"}, {"wav", "audio/wav"}, {"ogg", "audio/ogg"}}) {
var complete = service.toResponse("{\"data\":{\"status\":\"completed\",\"outputs\":[\"https://example.com/result."
+ sample[0] + "?token=test\"]}}", "audio");
assertEquals(sample[1], complete.getMimeType());
}
}
}