popiart-server/internal/server/models_test.go

package server

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"testing"
)

func TestHandleModelsLoadsDynamicCatalogFromNewAPI(t *testing.T) {
	var upstreamAuth string

	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v1/models" {
			http.NotFound(w, r)
			return
		}
		upstreamAuth = r.Header.Get("Authorization")
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{
			"success": true,
			"data": [
				{"id":"gpt-image-1","owned_by":"openai","supported_endpoint_types":["image-generation"]},
				{"id":"seedream-4-5-251128","owned_by":"custom","supported_endpoint_types":[]},
				{"id":"sora-2","owned_by":"openai","supported_endpoint_types":["openai-video"]},
				{"id":"viduq2","owned_by":"custom","supported_endpoint_types":[]},
				{"id":"doubao-seedance-2-0-260128","owned_by":"custom","supported_endpoint_types":[]},
				{"id":"custom-i2v-preview","owned_by":"custom","supported_endpoint_types":[]},
				{"id":"gpt-4.1","owned_by":"openai","supported_endpoint_types":["openai"]}
			]
		}`))
	}))
	defer upstream.Close()

	server, err := NewWithConfig(Config{
		NewAPIBaseURL: upstream.URL,
		SQLitePath:    filepath.Join(t.TempDir(), "popiart.db"),
		SkillhubDir:   makeEmptySkillhub(t),
		SessionSecret: "test-secret",
	})
	if err != nil {
		t.Fatalf("NewWithConfig: %v", err)
	}

	sessionToken, _, ok, err := server.store.createSession("sk-model-list-user")
	if err != nil {
		t.Fatalf("createSession: %v", err)
	}
	if !ok {
		t.Fatal("expected session creation to succeed")
	}

	req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	rec := httptest.NewRecorder()
	server.Handler().ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d: %s", rec.Code, rec.Body.String())
	}
	if upstreamAuth != "Bearer sk-model-list-user" {
		t.Fatalf("expected upstream auth header to use session upstream key, got %q", upstreamAuth)
	}

	var resp struct {
		OK   bool `json:"ok"`
		Data struct {
			Items []model `json:"items"`
		} `json:"data"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if !resp.OK {
		t.Fatalf("expected ok response, got %s", rec.Body.String())
	}
	if len(resp.Data.Items) != 6 {
		t.Fatalf("expected 6 supported media models, got %d: %s", len(resp.Data.Items), rec.Body.String())
	}

	byID := make(map[string]model, len(resp.Data.Items))
	for _, item := range resp.Data.Items {
		byID[item.ID] = item
	}
	if _, exists := byID["gpt-4.1"]; exists {
		t.Fatalf("expected unsupported text model to be filtered out: %+v", byID["gpt-4.1"])
	}

	imageModel, exists := byID["gpt-image-1"]
	if !exists {
		t.Fatal("expected gpt-image-1 in dynamic model list")
	}
	if imageModel.Type != "image" {
		t.Fatalf("expected gpt-image-1 type image, got %q", imageModel.Type)
	}
	if imageModel.Provider != "openai" {
		t.Fatalf("expected gpt-image-1 provider openai, got %q", imageModel.Provider)
	}
	if len(imageModel.Capabilities) != 2 || imageModel.Capabilities[0] != "text2image" || imageModel.Capabilities[1] != "img2img" {
		t.Fatalf("unexpected image capabilities: %#v", imageModel.Capabilities)
	}

	seedreamModel, exists := byID["seedream-4-5-251128"]
	if !exists {
		t.Fatal("expected seedream-4-5-251128 in dynamic model list")
	}
	if seedreamModel.Type != "image" {
		t.Fatalf("expected seedream-4-5-251128 type image, got %q", seedreamModel.Type)
	}
	if len(seedreamModel.Capabilities) != 2 || seedreamModel.Capabilities[0] != "text2image" || seedreamModel.Capabilities[1] != "img2img" {
		t.Fatalf("unexpected seedream capabilities: %#v", seedreamModel.Capabilities)
	}

	videoModel, exists := byID["sora-2"]
	if !exists {
		t.Fatal("expected sora-2 in dynamic model list")
	}
	if videoModel.Type != "video" {
		t.Fatalf("expected sora-2 type video, got %q", videoModel.Type)
	}
	if len(videoModel.Capabilities) != 2 || videoModel.Capabilities[0] != "text2video" || videoModel.Capabilities[1] != "image2video" {
		t.Fatalf("unexpected video capabilities: %#v", videoModel.Capabilities)
	}

	viduModel, exists := byID["viduq2"]
	if !exists {
		t.Fatal("expected viduq2 in dynamic model list")
	}
	if viduModel.Type != "video" {
		t.Fatalf("expected viduq2 type video, got %q", viduModel.Type)
	}
	if viduModel.Provider != "vidu" {
		t.Fatalf("expected viduq2 provider vidu, got %q", viduModel.Provider)
	}

	seedanceModel, exists := byID["doubao-seedance-2-0-260128"]
	if !exists {
		t.Fatal("expected doubao-seedance-2-0-260128 in dynamic model list")
	}
	if seedanceModel.Type != "video" {
		t.Fatalf("expected doubao-seedance-2-0-260128 type video, got %q", seedanceModel.Type)
	}
	if seedanceModel.Provider != "volcengine" {
		t.Fatalf("expected doubao-seedance-2-0-260128 provider volcengine, got %q", seedanceModel.Provider)
	}

	fallbackModel, exists := byID["custom-i2v-preview"]
	if !exists {
		t.Fatal("expected fallback-classified video model in dynamic model list")
	}
	if fallbackModel.Type != "video" {
		t.Fatalf("expected custom-i2v-preview fallback type video, got %q", fallbackModel.Type)
	}

	filterReq := httptest.NewRequest(http.MethodGet, "/v1/models?type=image&provider=openai", nil)
	filterReq.Header.Set("Authorization", "Bearer "+sessionToken)
	filterRec := httptest.NewRecorder()
	server.Handler().ServeHTTP(filterRec, filterReq)

	if filterRec.Code != http.StatusOK {
		t.Fatalf("expected filtered request 200, got %d: %s", filterRec.Code, filterRec.Body.String())
	}
	var filteredResp struct {
		Data struct {
			Items []model `json:"items"`
		} `json:"data"`
	}
	if err := json.Unmarshal(filterRec.Body.Bytes(), &filteredResp); err != nil {
		t.Fatalf("decode filtered response: %v", err)
	}
	if len(filteredResp.Data.Items) != 1 || filteredResp.Data.Items[0].ID != "gpt-image-1" {
		t.Fatalf("unexpected filtered items: %#v", filteredResp.Data.Items)
	}
}

func TestHandleModelsReturnsBadGatewayWhenUpstreamFails(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v1/models" {
			http.NotFound(w, r)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusBadGateway)
		_, _ = w.Write([]byte(`{"error":{"message":"upstream exploded"}}`))
	}))
	defer upstream.Close()

	server, err := NewWithConfig(Config{
		NewAPIBaseURL: upstream.URL,
		SQLitePath:    filepath.Join(t.TempDir(), "popiart.db"),
		SkillhubDir:   makeEmptySkillhub(t),
		SessionSecret: "test-secret",
	})
	if err != nil {
		t.Fatalf("NewWithConfig: %v", err)
	}

	sessionToken, _, ok, err := server.store.createSession("sk-model-list-user")
	if err != nil {
		t.Fatalf("createSession: %v", err)
	}
	if !ok {
		t.Fatal("expected session creation to succeed")
	}

	req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	rec := httptest.NewRecorder()
	server.Handler().ServeHTTP(rec, req)

	if rec.Code != http.StatusBadGateway {
		t.Fatalf("expected 502, got %d: %s", rec.Code, rec.Body.String())
	}

	var resp struct {
		OK    bool `json:"ok"`
		Error struct {
			Code    string `json:"code"`
			Message string `json:"message"`
		} `json:"error"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
		t.Fatalf("decode error response: %v", err)
	}
	if resp.OK {
		t.Fatalf("expected error response, got %s", rec.Body.String())
	}
	if resp.Error.Code != "MODEL_LIST_FAILED" {
		t.Fatalf("expected MODEL_LIST_FAILED, got %q", resp.Error.Code)
	}
}