feat: add markdown link extraction and HF URL normalization in ComfyModelAnalyzer and ModelSearchService
This commit is contained in:
Thomas Kippster
@@ -26,7 +26,7 @@ public class ComfyModelAnalyzerTest {
|
||||
.orElse(null);
|
||||
|
||||
assertNotNull(flux);
|
||||
assertEquals("checkpoints", flux.getType());
|
||||
assertEquals("unet", flux.getType());
|
||||
assertEquals("https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors", flux.getUrl());
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
package de.tki.comfymodels.service;
|
||||
|
||||
import com.sun.net.httpserver.HttpServer;
|
||||
import de.tki.comfymodels.domain.ModelInfo;
|
||||
import de.tki.comfymodels.service.impl.ConfigService;
|
||||
import de.tki.comfymodels.service.impl.DefaultDownloadManager;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import java.io.File;
|
||||
import java.io.IOException;
|
||||
import java.io.OutputStream;
|
||||
import java.lang.reflect.Field;
|
||||
import java.net.InetSocketAddress;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Collections;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.CopyOnWriteArrayList;
|
||||
import java.util.concurrent.CountDownLatch;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||
|
||||
public class DownloadStatusTest {
|
||||
|
||||
private DefaultDownloadManager downloadManager;
|
||||
private HttpServer server;
|
||||
private String serverUrl;
|
||||
private Path tempDir;
|
||||
|
||||
@BeforeEach
|
||||
public void setup() throws Exception {
|
||||
tempDir = Files.createTempDirectory("downloader_test");
|
||||
downloadManager = new DefaultDownloadManager();
|
||||
|
||||
// Use real ConfigService
|
||||
ConfigService configService = new ConfigService();
|
||||
|
||||
Field field = DefaultDownloadManager.class.getDeclaredField("configService");
|
||||
field.setAccessible(true);
|
||||
field.set(downloadManager, configService);
|
||||
|
||||
// Setup Local Test Server
|
||||
server = HttpServer.create(new InetSocketAddress(0), 0);
|
||||
server.createContext("/test.safetensors", exchange -> {
|
||||
byte[] data = new byte[1024 * 1024]; // 1MB Test Data
|
||||
exchange.getResponseHeaders().add("Content-Length", String.valueOf(data.length));
|
||||
exchange.sendResponseHeaders(200, data.length);
|
||||
try (OutputStream os = exchange.getResponseBody()) {
|
||||
// Sende Daten langsam, damit wir Pause/Stop testen können
|
||||
for (int i = 0; i < 10; i++) {
|
||||
os.write(data, i * (data.length / 10), data.length / 10);
|
||||
os.flush();
|
||||
Thread.sleep(100);
|
||||
}
|
||||
} catch (InterruptedException ignored) {}
|
||||
});
|
||||
server.start();
|
||||
serverUrl = "http://localhost:" + server.getAddress().getPort() + "/test.safetensors";
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
public void tearDown() throws IOException {
|
||||
server.stop(0);
|
||||
// Clean temp files
|
||||
Files.walk(tempDir).map(Path::toFile).forEach(File::delete);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testFullStatusLifecycle() throws InterruptedException {
|
||||
List<String> capturedStatus = new CopyOnWriteArrayList<>();
|
||||
CountDownLatch finishedLatch = new CountDownLatch(1);
|
||||
|
||||
ModelInfo info = new ModelInfo("checkpoints", "test_model.safetensors", serverUrl);
|
||||
List<ModelInfo> list = Collections.singletonList(info);
|
||||
boolean[] selected = {true};
|
||||
|
||||
downloadManager.startQueue(list, selected, tempDir.toString(),
|
||||
(idx, status) -> {
|
||||
capturedStatus.add(status);
|
||||
System.out.println("Status Update: " + status);
|
||||
},
|
||||
finishedLatch::countDown
|
||||
);
|
||||
|
||||
// 1. Prüfe auf Start/Download
|
||||
Thread.sleep(300);
|
||||
assertTrue(capturedStatus.stream().anyMatch(s -> s.contains("Downloading")), "Should show Downloading status");
|
||||
|
||||
// 2. Teste PAUSE
|
||||
System.out.println("--- Triggering PAUSE ---");
|
||||
downloadManager.togglePause();
|
||||
Thread.sleep(500);
|
||||
assertTrue(capturedStatus.contains("Paused"), "Should show Paused status");
|
||||
|
||||
// 3. Teste RESUME
|
||||
System.out.println("--- Triggering RESUME ---");
|
||||
downloadManager.togglePause();
|
||||
Thread.sleep(500);
|
||||
assertTrue(capturedStatus.contains("Resuming..."), "Should show Resuming status");
|
||||
|
||||
// 4. Teste STOP
|
||||
System.out.println("--- Triggering STOP ---");
|
||||
downloadManager.stop();
|
||||
|
||||
finishedLatch.await(5, TimeUnit.SECONDS);
|
||||
// Accept both "Stopped" or "Error: closed" (which is the result of closing the stream during stop)
|
||||
assertTrue(capturedStatus.stream().anyMatch(s -> s.contains("Stopped") || s.contains("Error: closed")),
|
||||
"Should show Stopped or closed status after stop(), but got: " + capturedStatus);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
package de.tki.comfymodels.service;
|
||||
|
||||
import de.tki.comfymodels.domain.ModelInfo;
|
||||
import de.tki.comfymodels.service.impl.ComfyModelAnalyzer;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
public class HiDreamReproductionTest {
|
||||
|
||||
@Test
|
||||
public void testHiDreamModelDetection() {
|
||||
String json = "{" +
|
||||
"\"id\": \"01d66ae9-78be-4a8d-b737-24eee5e1d447\"," +
|
||||
"\"nodes\": [" +
|
||||
" {\"id\": 55, \"type\": \"VAELoader\", \"widgets_values\": [\"ae.safetensors\"]}," +
|
||||
" {\"id\": 69, \"type\": \"UNETLoader\", \"widgets_values\": [\"hidream_i1_dev_bf16.safetensors\", \"default\"]}," +
|
||||
" {\"id\": 54, \"type\": \"QuadrupleCLIPLoader\", \"widgets_values\": [" +
|
||||
" \"clip_l_hidream.safetensors\", \"clip_g_hidream.safetensors\", " +
|
||||
" \"t5xxl_fp8_e4m3fn_scaled.safetensors\", \"llama_3.1_8b_instruct_fp8_scaled.safetensors\"" +
|
||||
" ]}," +
|
||||
" {\"id\": 72, \"type\": \"MarkdownNote\", \"widgets_values\": [\"## Official sampling settings\\n\\n### HiDream Full\\n* hidream_i1_full_fp16.safetensors\"]}" +
|
||||
"]" +
|
||||
"}";
|
||||
|
||||
ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer();
|
||||
List<ModelInfo> models = analyzer.analyze(json, "workflow.json");
|
||||
|
||||
// Wir erwarten 7 Modelle: 1 UNET, 4 CLIPs, 1 VAE (aus Loadern) UND 1 UNET (aus der MarkdownNote)
|
||||
assertEquals(7, models.size(), "Should detect 7 models (including one from MarkdownNote)");
|
||||
|
||||
boolean foundUnet = false;
|
||||
int clipsFound = 0;
|
||||
boolean foundVae = false;
|
||||
|
||||
for (ModelInfo m : models) {
|
||||
if (m.getName().equals("hidream_i1_dev_bf16.safetensors")) {
|
||||
assertEquals("unet", m.getType());
|
||||
foundUnet = true;
|
||||
} else if (m.getName().equals("ae.safetensors")) {
|
||||
assertEquals("vae", m.getType());
|
||||
foundVae = true;
|
||||
} else if (m.getName().contains("clip") || m.getName().contains("t5xxl") || m.getName().contains("llama")) {
|
||||
assertEquals("clip", m.getType(), "Type for " + m.getName() + " should be clip");
|
||||
clipsFound++;
|
||||
}
|
||||
}
|
||||
|
||||
assertTrue(foundUnet, "UNET from Loader not found");
|
||||
assertTrue(foundVae, "VAE not found");
|
||||
assertEquals(4, clipsFound, "Should find 4 CLIP/TextEncoder models");
|
||||
|
||||
// Überprüfe ob der Kontext "hidream" erkannt wurde
|
||||
boolean contextDetected = false;
|
||||
for (ModelInfo m : models) {
|
||||
System.out.println("Model: " + m.getName() + " | Popularity: " + m.getPopularity());
|
||||
if (m.getPopularity() != null && m.getPopularity().toLowerCase().contains("hidream")) {
|
||||
contextDetected = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
assertTrue(contextDetected, "Context 'hidream' should be detected from workflow content");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
package de.tki.comfymodels.service;
|
||||
|
||||
import de.tki.comfymodels.domain.ModelInfo;
|
||||
import de.tki.comfymodels.service.impl.ComfyModelAnalyzer;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
public class QualityAssuranceTest {
|
||||
|
||||
@Test
|
||||
public void testFalsePositiveFiltering() {
|
||||
ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer();
|
||||
|
||||
// Workflow mit "None" und "default" Werten in Widgets
|
||||
String json = "{" +
|
||||
"\"nodes\": [" +
|
||||
" {\"id\": 1, \"type\": \"CheckpointLoaderSimple\", \"widgets_values\": [\"None\", \"default\"]}," +
|
||||
" {\"id\": 2, \"type\": \"LoraLoader\", \"widgets_values\": [\"undefined\", \"test.safetensors\"]}" +
|
||||
"]" +
|
||||
"}";
|
||||
|
||||
List<ModelInfo> models = analyzer.analyze(json, "test.json");
|
||||
|
||||
// Es sollte nur "test.safetensors" gefunden werden
|
||||
assertEquals(1, models.size(), "Should only find one real model");
|
||||
assertEquals("test.safetensors", models.get(0).getName());
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testMalformedJsonHandling() {
|
||||
ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer();
|
||||
|
||||
// Kaputtes JSON
|
||||
String json = "{ \"nodes\": [ { \"id\": 1, \"type\": \"VAELoader\", \"widgets_values\": [\"ae.safetensors\""; // Fehlende Klammern
|
||||
|
||||
List<ModelInfo> models = analyzer.analyze(json, "broken.json");
|
||||
|
||||
// Sollte nicht abstürzen, sondern einfach 0 oder gefundene Modelle (je nach Parser-Robustheit) zurückgeben
|
||||
assertNotNull(models);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
package de.tki.comfymodels.service;
|
||||
|
||||
import de.tki.comfymodels.domain.ModelInfo;
|
||||
import de.tki.comfymodels.service.impl.ComfyModelAnalyzer;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import java.util.List;
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
|
||||
public class ReproductionTest {
|
||||
|
||||
private final ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer();
|
||||
|
||||
@Test
|
||||
public void testDeduplicationAndUrlFiltering() {
|
||||
String json = "{\n" +
|
||||
" \"nodes\": [\n" +
|
||||
" {\n" +
|
||||
" \"type\": \"MarkdownNote\",\n" +
|
||||
" \"widgets_values\": [\n" +
|
||||
" \"- [qwen_3_4b.safetensors](https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/split_files/text_encoders/qwen_3_4b.safetensors)\"\n" +
|
||||
" ]\n" +
|
||||
" },\n" +
|
||||
" {\n" +
|
||||
" \"type\": \"CLIPLoader\",\n" +
|
||||
" \"properties\": {\n" +
|
||||
" \"models\": [\n" +
|
||||
" {\n" +
|
||||
" \"name\": \"qwen_3_4b.safetensors\",\n" +
|
||||
" \"url\": \"https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/split_files/text_encoders/qwen_3_4b.safetensors\",\n" +
|
||||
" \"directory\": \"text_encoders\"\n" +
|
||||
" }\n" +
|
||||
" ]\n" +
|
||||
" }\n" +
|
||||
" }\n" +
|
||||
" ]\n" +
|
||||
"}";
|
||||
|
||||
List<ModelInfo> models = analyzer.analyze(json, "test.json");
|
||||
|
||||
// Erwartung: Nur EIN Eintrag für qwen_3_4b.safetensors
|
||||
// Der Name darf NICHT die URL sein.
|
||||
long count = models.stream()
|
||||
.filter(m -> m.getName().toLowerCase().contains("qwen"))
|
||||
.count();
|
||||
|
||||
for (ModelInfo m : models) {
|
||||
System.out.println("Found: " + m.getName() + " [" + m.getType() + "] URL: " + m.getUrl());
|
||||
}
|
||||
|
||||
assertEquals(1, count, "Should only find one entry for qwen, not a second one with the URL as name");
|
||||
assertEquals("qwen_3_4b.safetensors", models.get(0).getName());
|
||||
assertEquals("text_encoders", models.get(0).getType());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
package de.tki.comfymodels.service;
|
||||
|
||||
import de.tki.comfymodels.domain.ModelInfo;
|
||||
import de.tki.comfymodels.service.impl.ConfigService;
|
||||
import de.tki.comfymodels.service.impl.DefaultDownloadManager;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.test.util.ReflectionTestUtils;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.Arrays;
|
||||
import java.util.Collections;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.concurrent.CountDownLatch;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||
|
||||
public class DownloadSelectionTest {
|
||||
|
||||
@Test
|
||||
public void testOnlySelectedItemsAreProcessed() throws InterruptedException {
|
||||
DefaultDownloadManager downloadManager = new DefaultDownloadManager();
|
||||
|
||||
// Mocking ConfigService to avoid NPE
|
||||
ConfigService configService = new ConfigService();
|
||||
ReflectionTestUtils.setField(downloadManager, "configService", configService);
|
||||
|
||||
List<ModelInfo> models = new ArrayList<>();
|
||||
models.add(new ModelInfo("checkpoints", "model1.safetensors", "http://example.com/1"));
|
||||
models.add(new ModelInfo("checkpoints", "model2.safetensors", "http://example.com/2"));
|
||||
models.add(new ModelInfo("checkpoints", "model3.safetensors", "http://example.com/3"));
|
||||
|
||||
// Only first and third are selected
|
||||
boolean[] selected = {true, false, true};
|
||||
|
||||
ConcurrentHashMap<Integer, String> statusMap = new ConcurrentHashMap<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
downloadManager.startQueue(models, selected, "target/test_downloads", (idx, status) -> {
|
||||
statusMap.put(idx, status);
|
||||
}, () -> {
|
||||
latch.countDown();
|
||||
});
|
||||
|
||||
// Wait for completion
|
||||
assertTrue(latch.await(5, TimeUnit.SECONDS), "Download queue did not finish in time");
|
||||
|
||||
// Model 2 MUST be skipped
|
||||
assertEquals("Skipped (Not Selected)", statusMap.get(1), "Model 2 should have been skipped");
|
||||
|
||||
// Model 1 and 3 should have been attempted (likely Error because URL is fake, but NOT Skipped (Not Selected))
|
||||
assertTrue(statusMap.containsKey(0));
|
||||
assertTrue(!statusMap.get(0).equals("Skipped (Not Selected)"));
|
||||
assertTrue(statusMap.containsKey(2));
|
||||
assertTrue(!statusMap.get(2).equals("Skipped (Not Selected)"));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testDynamicUncheckWhilePaused() throws InterruptedException {
|
||||
DefaultDownloadManager downloadManager = new DefaultDownloadManager();
|
||||
ConfigService configService = new ConfigService();
|
||||
ReflectionTestUtils.setField(downloadManager, "configService", configService);
|
||||
|
||||
List<ModelInfo> models = new ArrayList<>();
|
||||
// Use a dummy model
|
||||
models.add(new ModelInfo("checkpoints", "model1.safetensors", "http://invalid-url-for-test.com/file"));
|
||||
|
||||
boolean[] selected = {true};
|
||||
ConcurrentHashMap<Integer, String> statusMap = new ConcurrentHashMap<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
// Start in PAUSED state
|
||||
downloadManager.togglePause();
|
||||
|
||||
downloadManager.startQueue(models, selected, "target/test_downloads", (idx, status) -> {
|
||||
statusMap.put(idx, status);
|
||||
}, () -> {
|
||||
latch.countDown();
|
||||
});
|
||||
|
||||
// Give it a moment to reach the pause loop in startQueue or downloadWithResume
|
||||
Thread.sleep(300);
|
||||
|
||||
// At this point, it should be either "Idle" (not yet started) or "Paused"
|
||||
// But the important thing is: We uncheck it NOW while it's waiting
|
||||
downloadManager.updateSelection(new boolean[]{false});
|
||||
|
||||
// Resume the manager
|
||||
downloadManager.togglePause();
|
||||
|
||||
// Wait for completion (the thread should now see it's unchecked and finish)
|
||||
assertTrue(latch.await(5, TimeUnit.SECONDS), "Download queue did not finish in time");
|
||||
|
||||
// The status MUST indicate it was skipped or errored out (as long as it didn't finish)
|
||||
String finalStatus = statusMap.get(0);
|
||||
assertTrue(finalStatus != null && (finalStatus.contains("Skipped") || finalStatus.contains("ConnectException")),
|
||||
"Expected 'Skipped' or 'ConnectException' status but got: " + finalStatus);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package de.tki.comfymodels.service;
|
||||
|
||||
import de.tki.comfymodels.domain.ModelInfo;
|
||||
import de.tki.comfymodels.service.impl.ComfyModelAnalyzer;
|
||||
import org.json.JSONObject;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
|
||||
public class ModelTypePathTest {
|
||||
|
||||
@Test
|
||||
public void testModelTypeMapping() {
|
||||
ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer();
|
||||
|
||||
// Mock JSON representing different nodes
|
||||
String json = "{" +
|
||||
"\"1\": {\"class_type\": \"LoraLoader\", \"widgets_values\": [\"my_lora.safetensors\"]}," +
|
||||
"\"2\": {\"class_type\": \"ControlNetApply\", \"widgets_values\": [\"cn_model.pth\"]}," +
|
||||
"\"3\": {\"class_type\": \"CLIPLoader\", \"widgets_values\": [\"t5xxl.safetensors\"]}," +
|
||||
"\"4\": {\"class_type\": \"VAELoader\", \"widgets_values\": [\"vae_model.safetensors\"]}," +
|
||||
"\"5\": {\"class_type\": \"UNETLoader\", \"widgets_values\": [\"diffusion_model.safetensors\"]}," +
|
||||
"\"6\": {\"class_type\": \"GLIGENLoader\", \"widgets_values\": [\"gligen_box.safetensors\"]}," +
|
||||
"\"7\": {\"class_type\": \"CLIPVisionLoader\", \"widgets_values\": [\"clip_vision.safetensors\"]}," +
|
||||
"\"8\": {\"class_type\": \"AudioEncoderLoader\", \"widgets_values\": [\"audio.safetensors\"]}" +
|
||||
"}";
|
||||
|
||||
List<ModelInfo> results = analyzer.analyze(json, "test_workflow.json");
|
||||
results.forEach(r -> System.out.println("Found: " + r.getName() + " -> " + r.getType()));
|
||||
|
||||
assertPath(results, "my_lora.safetensors", "loras");
|
||||
assertPath(results, "cn_model.pth", "controlnet");
|
||||
assertPath(results, "t5xxl.safetensors", "clip");
|
||||
assertPath(results, "vae_model.safetensors", "vae");
|
||||
assertPath(results, "diffusion_model.safetensors", "unet");
|
||||
assertPath(results, "gligen_box.safetensors", "gligen");
|
||||
assertPath(results, "clip_vision.safetensors", "clip_vision");
|
||||
assertPath(results, "audio.safetensors", "audio_encoders");
|
||||
}
|
||||
|
||||
private void assertPath(List<ModelInfo> results, String name, String expectedType) {
|
||||
for (ModelInfo info : results) {
|
||||
if (info.getName().equals(name)) {
|
||||
assertEquals(expectedType, info.getType(), "Type mismatch for " + name);
|
||||
return;
|
||||
}
|
||||
}
|
||||
throw new AssertionError("Model " + name + " not found in results");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package de.tki.comfymodels.service;
|
||||
|
||||
import de.tki.comfymodels.domain.ModelInfo;
|
||||
import de.tki.comfymodels.service.impl.ComfyModelAnalyzer;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import java.util.List;
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
|
||||
public class QwenReproductionTest {
|
||||
|
||||
@Test
|
||||
public void testQwenMarkdownModelDetection() {
|
||||
ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer();
|
||||
// Minimaler Workflow mit der MarkdownNote des Nutzers
|
||||
String workflow = "{\n" +
|
||||
" \"nodes\": [\n" +
|
||||
" {\n" +
|
||||
" \"id\": 99,\n" +
|
||||
" \"type\": \"MarkdownNote\",\n" +
|
||||
" \"widgets_values\": [\n" +
|
||||
" \"## Model links\\n\\n**LoRA**\\n\\n- [Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors](https://huggingface.co/lightx2v/Qwen-Image-Lightning/resolve/main/Qwen-Image-Edit-2509/Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors)\\n\\n📂 loras/\\n├── 📂 loras/\\n│ └── Qwen-Image-Lightning-4steps-V1.0.safetensors\"\n" +
|
||||
" ]\n" +
|
||||
" }\n" +
|
||||
" ]\n" +
|
||||
"}";
|
||||
|
||||
List<ModelInfo> models = analyzer.analyze(workflow, "test.json");
|
||||
|
||||
// Wir prüfen, ob beide Namen gefunden werden
|
||||
boolean foundLightning = models.stream().anyMatch(m -> m.getName().equals("Qwen-Image-Lightning-4steps-V1.0.safetensors"));
|
||||
boolean foundEdit = models.stream().anyMatch(m -> m.getName().equals("Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors"));
|
||||
|
||||
assertTrue(foundEdit, "Modell aus dem Link sollte gefunden werden");
|
||||
assertTrue(foundLightning, "Modell aus dem Text-Diagramm sollte gefunden werden");
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testMarkdownLinkUrlExtraction() {
|
||||
ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer();
|
||||
String workflow = "{\n" +
|
||||
" \"nodes\": [\n" +
|
||||
" {\n" +
|
||||
" \"type\": \"MarkdownNote\",\n" +
|
||||
" \"widgets_values\": [\n" +
|
||||
" \"- [Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors](https://huggingface.co/lightx2v/Qwen-Image-Lightning/resolve/main/Qwen-Image-Edit-2509/Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors)\"\n" +
|
||||
" ]\n" +
|
||||
" }\n" +
|
||||
" ]\n" +
|
||||
"}";
|
||||
|
||||
List<ModelInfo> models = analyzer.analyze(workflow, "test.json");
|
||||
|
||||
ModelInfo editModel = models.stream()
|
||||
.filter(m -> m.getName().equals("Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors"))
|
||||
.findFirst().orElse(null);
|
||||
|
||||
assertNotNull(editModel, "Modell sollte gefunden werden");
|
||||
assertEquals("https://huggingface.co/lightx2v/Qwen-Image-Lightning/resolve/main/Qwen-Image-Edit-2509/Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors",
|
||||
editModel.getUrl(), "URL sollte aus dem Markdown-Link extrahiert werden");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,220 @@
|
||||
package de.tki.comfymodels.service;
|
||||
|
||||
import com.sun.net.httpserver.HttpServer;
|
||||
import de.tki.comfymodels.domain.ModelInfo;
|
||||
import de.tki.comfymodels.service.impl.*;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.test.util.ReflectionTestUtils;
|
||||
|
||||
import java.io.File;
|
||||
import java.io.IOException;
|
||||
import java.io.OutputStream;
|
||||
import java.net.InetSocketAddress;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.util.Collections;
|
||||
import java.util.Comparator;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.CopyOnWriteArrayList;
|
||||
import java.util.concurrent.CountDownLatch;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||
|
||||
public class StatusTransitionTest {
|
||||
|
||||
private DefaultDownloadManager downloadManager;
|
||||
private HttpServer server;
|
||||
private String serverUrl;
|
||||
private Path tempDir;
|
||||
private ConfigService configService;
|
||||
|
||||
@BeforeEach
|
||||
public void setup() throws Exception {
|
||||
tempDir = Files.createTempDirectory("transition_test");
|
||||
downloadManager = new DefaultDownloadManager();
|
||||
configService = new ConfigService();
|
||||
ReflectionTestUtils.setField(downloadManager, "configService", configService);
|
||||
|
||||
// Setup Local Test Server
|
||||
server = HttpServer.create(new InetSocketAddress(0), 0);
|
||||
server.createContext("/test.safetensors", exchange -> {
|
||||
if ("HEAD".equalsIgnoreCase(exchange.getRequestMethod())) {
|
||||
exchange.getResponseHeaders().add("Content-Length", "10000");
|
||||
exchange.sendResponseHeaders(200, -1);
|
||||
return;
|
||||
}
|
||||
byte[] data = new byte[10000];
|
||||
exchange.sendResponseHeaders(200, data.length);
|
||||
try (OutputStream os = exchange.getResponseBody()) {
|
||||
for (int i = 0; i < 5; i++) {
|
||||
os.write(data, i * 2000, 2000);
|
||||
os.flush();
|
||||
Thread.sleep(100);
|
||||
}
|
||||
} catch (Exception ignored) {}
|
||||
});
|
||||
|
||||
server.createContext("/auth-fail", exchange -> {
|
||||
exchange.sendResponseHeaders(401, -1);
|
||||
});
|
||||
|
||||
server.start();
|
||||
serverUrl = "http://localhost:" + server.getAddress().getPort();
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
public void tearDown() throws IOException {
|
||||
server.stop(0);
|
||||
if (Files.exists(tempDir)) {
|
||||
Files.walk(tempDir)
|
||||
.sorted(Comparator.reverseOrder())
|
||||
.map(Path::toFile)
|
||||
.forEach(File::delete);
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testTransition_Idle_To_Downloading_To_Finished() throws InterruptedException {
|
||||
List<String> statuses = new CopyOnWriteArrayList<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
ModelInfo info = new ModelInfo("checkpoints", "test.safetensors", serverUrl + "/test.safetensors");
|
||||
downloadManager.startQueue(Collections.singletonList(info), new boolean[]{true}, tempDir.toString(),
|
||||
(idx, s) -> statuses.add(s), latch::countDown);
|
||||
|
||||
assertTrue(latch.await(5, TimeUnit.SECONDS));
|
||||
assertTrue(statuses.stream().anyMatch(s -> s.contains("Downloading")), "Statuses: " + statuses);
|
||||
assertTrue(statuses.contains("✅ Finished"), "Statuses: " + statuses);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testTransition_Idle_To_Searching_To_NotFound() throws InterruptedException {
|
||||
ModelSearchService searchService = new ModelSearchService();
|
||||
ConfigService cfg = new ConfigService();
|
||||
GeminiAIService gemini = new GeminiAIService();
|
||||
ModelListService list = new ModelListService();
|
||||
|
||||
ReflectionTestUtils.setField(searchService, "configService", cfg);
|
||||
ReflectionTestUtils.setField(searchService, "geminiService", gemini);
|
||||
ReflectionTestUtils.setField(searchService, "modelListService", list);
|
||||
ReflectionTestUtils.setField(gemini, "configService", cfg);
|
||||
|
||||
List<String> statuses = new CopyOnWriteArrayList<>();
|
||||
ModelInfo info = new ModelInfo("checkpoints", "non_existent_model_xyz.safetensors", "MISSING");
|
||||
|
||||
searchService.searchOnline(Collections.singletonList(info), new boolean[]{true}, "context", "file.json",
|
||||
(idx, s) -> statuses.add(s),
|
||||
(idx, found) -> {},
|
||||
() -> {}
|
||||
);
|
||||
|
||||
// Wait a bit for search to progress (it will fail anyway as no API key is set)
|
||||
Thread.sleep(1500);
|
||||
assertTrue(statuses.stream().anyMatch(s -> s.contains("Searching") || s.contains("Scouting")), "Statuses: " + statuses);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testTransition_Idle_To_SkippedNotSelected() throws InterruptedException {
|
||||
List<String> statuses = new CopyOnWriteArrayList<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
ModelInfo info = new ModelInfo("checkpoints", "skip_test.safetensors", serverUrl + "/test.safetensors");
|
||||
downloadManager.startQueue(Collections.singletonList(info), new boolean[]{false}, tempDir.toString(),
|
||||
(idx, s) -> statuses.add(s), latch::countDown);
|
||||
|
||||
assertTrue(latch.await(2, TimeUnit.SECONDS));
|
||||
assertTrue(statuses.contains("Skipped (Not Selected)"));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testTransition_Downloading_To_Paused_To_Resuming() throws InterruptedException {
|
||||
List<String> statuses = new CopyOnWriteArrayList<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
ModelInfo info = new ModelInfo("checkpoints", "pause_test.safetensors", serverUrl + "/test.safetensors");
|
||||
downloadManager.startQueue(Collections.singletonList(info), new boolean[]{true}, tempDir.toString(),
|
||||
(idx, s) -> statuses.add(s), latch::countDown);
|
||||
|
||||
Thread.sleep(300);
|
||||
downloadManager.togglePause();
|
||||
Thread.sleep(500);
|
||||
assertTrue(statuses.contains("Paused"), "Statuses: " + statuses);
|
||||
|
||||
downloadManager.togglePause(); // Resume
|
||||
Thread.sleep(500);
|
||||
assertTrue(statuses.contains("Resuming..."), "Statuses: " + statuses);
|
||||
|
||||
downloadManager.stop();
|
||||
latch.await(2, TimeUnit.SECONDS);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testTransition_Paused_To_SkippedUnchecked() throws InterruptedException {
|
||||
List<String> statuses = new CopyOnWriteArrayList<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
ModelInfo info = new ModelInfo("checkpoints", "uncheck_test.safetensors", serverUrl + "/test.safetensors");
|
||||
downloadManager.startQueue(Collections.singletonList(info), new boolean[]{true}, tempDir.toString(),
|
||||
(idx, s) -> statuses.add(s), latch::countDown);
|
||||
|
||||
Thread.sleep(300);
|
||||
downloadManager.togglePause();
|
||||
Thread.sleep(300);
|
||||
|
||||
downloadManager.updateSelection(new boolean[]{false});
|
||||
downloadManager.togglePause(); // Resume triggers check
|
||||
|
||||
assertTrue(latch.await(2, TimeUnit.SECONDS));
|
||||
assertTrue(statuses.contains("Skipped (Unchecked)"), "Statuses: " + statuses);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testTransition_Downloading_To_Stopped() throws InterruptedException {
|
||||
List<String> statuses = new CopyOnWriteArrayList<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
ModelInfo info = new ModelInfo("checkpoints", "stop_test.safetensors", serverUrl + "/test.safetensors");
|
||||
downloadManager.startQueue(Collections.singletonList(info), new boolean[]{true}, tempDir.toString(),
|
||||
(idx, s) -> statuses.add(s), latch::countDown);
|
||||
|
||||
Thread.sleep(300);
|
||||
downloadManager.stop();
|
||||
|
||||
assertTrue(latch.await(2, TimeUnit.SECONDS));
|
||||
assertTrue(statuses.contains("Stopped"), "Statuses: " + statuses);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testTransition_Resuming_To_AuthRequired() throws InterruptedException {
|
||||
List<String> statuses = new CopyOnWriteArrayList<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
ModelInfo info = new ModelInfo("checkpoints", "auth.safetensors", serverUrl + "/auth-fail");
|
||||
|
||||
downloadManager.togglePause();
|
||||
downloadManager.startQueue(Collections.singletonList(info), new boolean[]{true}, tempDir.toString(),
|
||||
(idx, s) -> statuses.add(s), latch::countDown);
|
||||
|
||||
Thread.sleep(300);
|
||||
downloadManager.togglePause(); // Resume -> will hit 401
|
||||
|
||||
assertTrue(latch.await(2, TimeUnit.SECONDS));
|
||||
assertTrue(statuses.contains("❌ Auth Required (Token?)"), "Statuses: " + statuses);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testTransition_Downloading_To_Error() throws InterruptedException {
|
||||
List<String> statuses = new CopyOnWriteArrayList<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
ModelInfo info = new ModelInfo("checkpoints", "error.safetensors", "http://localhost:1");
|
||||
downloadManager.startQueue(Collections.singletonList(info), new boolean[]{true}, tempDir.toString(),
|
||||
(idx, s) -> statuses.add(s), latch::countDown);
|
||||
|
||||
assertTrue(latch.await(5, TimeUnit.SECONDS));
|
||||
assertTrue(statuses.stream().anyMatch(s -> s.startsWith("Error:")), "Statuses: " + statuses);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
package de.tki.comfymodels.service;
|
||||
|
||||
import de.tki.comfymodels.service.impl.ConfigService;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.io.TempDir;
|
||||
import java.io.File;
|
||||
import java.nio.file.Path;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
|
||||
public class VaultIntegrationTest {
|
||||
|
||||
@TempDir
|
||||
Path tempDir;
|
||||
|
||||
@Test
|
||||
public void testVaultSaveAndLoad() throws Exception {
|
||||
// Setup ConfigService with temp directory
|
||||
ConfigService configService = new ConfigService() {
|
||||
@Override
|
||||
public String getAppDataPath() {
|
||||
return tempDir.toString();
|
||||
}
|
||||
};
|
||||
|
||||
String password = "1234";
|
||||
String testKey = "4711";
|
||||
|
||||
// 1. Initial unlock (creates new vault)
|
||||
configService.unlock(password);
|
||||
|
||||
// 2. Set key and save
|
||||
configService.setGeminiApiKey(testKey);
|
||||
configService.save();
|
||||
|
||||
// 3. Create NEW instance to simulate restart
|
||||
ConfigService newConfigService = new ConfigService() {
|
||||
@Override
|
||||
public String getAppDataPath() {
|
||||
return tempDir.toString();
|
||||
}
|
||||
};
|
||||
|
||||
// 4. Unlock with correct password
|
||||
newConfigService.unlock(password);
|
||||
|
||||
// 5. Verify key
|
||||
assertEquals(testKey, newConfigService.getGeminiApiKey(), "Key should be 4711 after loading");
|
||||
|
||||
// 6. Verify wrong password fails
|
||||
ConfigService wrongPassService = new ConfigService() {
|
||||
@Override
|
||||
public String getAppDataPath() {
|
||||
return tempDir.toString();
|
||||
}
|
||||
};
|
||||
|
||||
assertThrows(Exception.class, () -> {
|
||||
wrongPassService.unlock("wrong_password");
|
||||
}, "Unlock should fail with wrong password");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user