Atlas/internal/provider/model_schema_test.go

package provider

import (
	"os"
	"testing"
)

func TestParseReplicateModelDetails(t *testing.T) {
	body, err := os.ReadFile("testdata/models_sample.json")
	if err != nil {
		t.Skip("sample file missing")
	}
	models := parseReplicateModelDetails(body)
	if len(models) != 2 {
		t.Fatalf("expected 2 models, got %d", len(models))
	}
	if models[0].Kind != ModelKindImageToImage {
		t.Fatalf("edit model kind: %s", models[0].Kind)
	}
	if models[1].Kind != ModelKindTextToImage {
		t.Fatalf("text model kind: %s", models[1].Kind)
	}
	if models[0].ModelName != "qwen-image-edit" {
		t.Fatalf("edit model name: %s", models[0].ModelName)
	}
	if models[0].PredictionVersion() != "comfyui/qwen-image-edit" {
		t.Fatalf("edit prediction version: %s", models[0].PredictionVersion())
	}
}

func TestInferModelKind(t *testing.T) {
	schema := []InputField{{Name: "prompt", Required: true}, {Name: "image", Required: true, Widget: "image"}}
	if inferModelKind(schema) != ModelKindImageToImage {
		t.Fatal("expected image_to_image")
	}
}

func TestValidateModelInput(t *testing.T) {
	schema := []InputField{{Name: "prompt", Required: true}}
	if err := ValidateModelInput(schema, map[string]any{"prompt": "hi"}); err != nil {
		t.Fatal(err)
	}
	if err := ValidateModelInput(schema, map[string]any{}); err == nil {
		t.Fatal("expected error")
	}
}