diff --git a/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/demo/wss/Main.java b/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/demo/wss/Main.java
deleted file mode 100644
index 25a332702..000000000
--- a/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/midjourney/demo/wss/Main.java
+++ /dev/null
@@ -1,39 +0,0 @@
-package cn.iocoder.yudao.framework.ai.midjourney.demo.wss;
-
-import cn.iocoder.yudao.framework.ai.midjourney.demo.jad.DiscordAccount;
-import cn.iocoder.yudao.framework.ai.midjourney.demo.wss.user.SpringUserWebSocketStarter;
-import cn.iocoder.yudao.framework.ai.midjourney.demo.wss.user.UserMessageListener;
-
-import java.util.Scanner;
-
-/**
- * author: fansili
- * time: 2024/4/3 16:40
- */
-public class Main {
-
-    public static void main(String[] args) {
-        String token =  "NTY5MDY4NDAxNzEyOTU1Mzky.G4-Fu0.MzD-7ll-ElbXTTgDPHF-WS_UyhMAfbKN3WyyBc";
-
-
-        DiscordHelper discordHelper = new DiscordHelper();
-        DiscordAccount discordAccount = new DiscordAccount();
-        discordAccount.setUserToken(token);
-        discordAccount.setGuildId("1221445697157468200");
-        discordAccount.setChannelId("1221445862962630706");
-
-
-        var messageListener = new UserMessageListener();
-        var webSocketStarter = new SpringUserWebSocketStarter(discordHelper.getWss(), null, discordAccount, messageListener);
-
-        try {
-            webSocketStarter.start();
-        } catch (Exception e) {
-            throw new RuntimeException(e);
-        }
-
-
-        Scanner scanner = new Scanner(System.in);
-        scanner.nextLine();
-    }
-}
diff --git a/yudao-module-ai/yudao-spring-boot-starter-ai/src/test/java/cn/iocoder/yudao/framework/ai/image/MjWebSocketTests.java b/yudao-module-ai/yudao-spring-boot-starter-ai/src/test/java/cn/iocoder/yudao/framework/ai/image/MjWebSocketTests.java
new file mode 100644
index 000000000..177b7c733
--- /dev/null
+++ b/yudao-module-ai/yudao-spring-boot-starter-ai/src/test/java/cn/iocoder/yudao/framework/ai/image/MjWebSocketTests.java
@@ -0,0 +1,51 @@
+package cn.iocoder.yudao.framework.ai.image;
+
+import cn.hutool.core.io.FileUtil;
+import cn.iocoder.yudao.framework.ai.midjourney.MidjourneyConfig;
+import cn.iocoder.yudao.framework.ai.midjourney.demo.wss.user.MjMessageListener;
+import cn.iocoder.yudao.framework.ai.midjourney.demo.wss.user.MjWebSocketStarter;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.io.File;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Scanner;
+
+/**
+ * author: fansili
+ * time: 2024/4/3 16:40
+ */
+public class MjWebSocketTests {
+
+    private MidjourneyConfig midjourneyConfig;
+
+    @Before
+    public void setup() {
+        String token =  "OTcwNDc3NzQxMjUyMTY5NzI4.GJcVxa.VrzMii8dsHOJAPZn4Mw8GuEo7_nIUJij9JIHD4";
+        Map<String, String> requestTemplates = new HashMap<>();
+        List<File> files = FileUtil.loopFiles("/Users/fansili/projects/github/ruoyi-vue-pro/yudao-module-ai/yudao-spring-boot-starter-ai/src/main/resources/http-body");
+        for (File file : files) {
+            requestTemplates.put(file.getName().replace(".json", ""), FileUtil.readUtf8String(file));
+        }
+        this.midjourneyConfig = new MidjourneyConfig(token, "1225414986084388926", "1225414986587832385", requestTemplates);
+    }
+
+    @Test
+    public void startSocketTest() {
+        String wssUrl = "wss://gateway.discord.gg";
+        var messageListener = new MjMessageListener(midjourneyConfig);
+        var webSocketStarter = new MjWebSocketStarter(wssUrl, null, midjourneyConfig, messageListener);
+
+        try {
+            webSocketStarter.start();
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+
+
+        Scanner scanner = new Scanner(System.in);
+        scanner.nextLine();
+    }
+}