popiart-server/internal/server/artifact_upload_test.go

package server

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"testing"
)

func TestArtifactUploadCreatesReadableArtifactForSourceArtifactID(t *testing.T) {
	cfg := Config{
		SQLitePath:    filepath.Join(t.TempDir(), "popiart.db"),
		SkillhubDir:   makeEmptySkillhub(t),
		SessionSecret: "test-secret",
	}
	server, err := NewWithConfig(cfg)
	if err != nil {
		t.Fatalf("NewWithConfig: %v", err)
	}
	srv := httptest.NewServer(server.Handler())
	defer srv.Close()
	server.cfg.PublicBaseURL = srv.URL

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

	imageBytes := tinyPNG(t)
	var body bytes.Buffer
	writer := multipart.NewWriter(&body)
	if err := writer.WriteField("role", "source"); err != nil {
		t.Fatalf("write role field: %v", err)
	}
	if err := writer.WriteField("metadata_json", `{"origin":"chat-upload"}`); err != nil {
		t.Fatalf("write metadata field: %v", err)
	}
	part, err := writer.CreateFormFile("file", "chat-upload.png")
	if err != nil {
		t.Fatalf("create form file: %v", err)
	}
	if _, err := part.Write(imageBytes); err != nil {
		t.Fatalf("write form file: %v", err)
	}
	if err := writer.Close(); err != nil {
		t.Fatalf("close multipart writer: %v", err)
	}

	req := httptest.NewRequest(http.MethodPost, "/v1/artifacts/upload", &body)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	req.Header.Set("Content-Type", writer.FormDataContentType())
	rec := httptest.NewRecorder()
	server.Handler().ServeHTTP(rec, req)

	if rec.Code != http.StatusCreated {
		t.Fatalf("expected status 201, got %d body=%s", rec.Code, rec.Body.String())
	}

	var envelope struct {
		OK   bool     `json:"ok"`
		Data artifact `json:"data"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &envelope); err != nil {
		t.Fatalf("decode upload response: %v body=%s", err, rec.Body.String())
	}
	if !envelope.OK {
		t.Fatalf("expected ok upload response, got %s", rec.Body.String())
	}
	if envelope.Data.ID == "" {
		t.Fatalf("expected artifact id, got %#v", envelope.Data)
	}
	if envelope.Data.MediaID == "" {
		t.Fatalf("expected media id, got %#v", envelope.Data)
	}
	if envelope.Data.ContentType != "image/png" {
		t.Fatalf("expected image/png, got %q", envelope.Data.ContentType)
	}
	if envelope.Data.URL == "" {
		t.Fatalf("expected stable artifact url, got %#v", envelope.Data)
	}
	if envelope.Data.StorageStatus != "ready" {
		t.Fatalf("expected storage status ready, got %#v", envelope.Data.StorageStatus)
	}
	if envelope.Data.SourceSkillID != "popiskill-artifact-upload-local-v1" {
		t.Fatalf("expected source skill id, got %#v", envelope.Data.SourceSkillID)
	}
	if envelope.Data.SourceRouteKey != "artifact.upload" {
		t.Fatalf("expected source route key, got %#v", envelope.Data.SourceRouteKey)
	}
	if envelope.Data.SourceInput["role"] != "source" {
		t.Fatalf("expected source input role, got %#v", envelope.Data.SourceInput)
	}

	contentResp, err := http.Get(envelope.Data.URL)
	if err != nil {
		t.Fatalf("GET media content: %v", err)
	}
	defer contentResp.Body.Close()
	if contentResp.StatusCode != http.StatusOK {
		t.Fatalf("expected media content 200, got %d", contentResp.StatusCode)
	}
	streamed, err := io.ReadAll(contentResp.Body)
	if err != nil {
		t.Fatalf("read media content: %v", err)
	}
	if !bytes.Equal(streamed, imageBytes) {
		t.Fatal("expected stable media url to serve uploaded bytes")
	}

	editRef, err := server.resolveImageToImageReference(context.Background(), &job{
		UserID:      current.User.ID,
		SessionID:   current.Token,
		UpstreamKey: current.UpstreamKey,
		SkillID:     "popiskill-image-img2img-basic-v1",
	}, map[string]any{
		"source_artifact_id": envelope.Data.ID,
	})
	if err != nil {
		t.Fatalf("resolveImageToImageReference: %v", err)
	}
	if editRef.ContentType != "image/png" {
		t.Fatalf("expected resolved content type image/png, got %q", editRef.ContentType)
	}
	if !bytes.Equal(editRef.Content, imageBytes) {
		t.Fatal("expected resolved reference content to match uploaded bytes")
	}

	unsignedResp, err := http.Get(strings.Split(envelope.Data.URL, "?")[0])
	if err != nil {
		t.Fatalf("GET unsigned artifact media content: %v", err)
	}
	defer unsignedResp.Body.Close()
	if unsignedResp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("expected unsigned artifact media content 401, got %d", unsignedResp.StatusCode)
	}

	repo, ok := server.store.artifacts.(*sqliteRepository)
	if !ok {
		t.Fatalf("expected sqliteRepository, got %T", server.store.artifacts)
	}
	if _, err := repo.db.Exec(`UPDATE jobs SET result_refs_json = NULL WHERE job_id = ?`, envelope.Data.JobID); err != nil {
		t.Fatalf("clear job result refs: %v", err)
	}

	getReq := httptest.NewRequest(http.MethodGet, "/v1/artifacts/"+envelope.Data.ID, nil)
	getReq.Header.Set("Authorization", "Bearer "+sessionToken)
	getRec := httptest.NewRecorder()
	server.Handler().ServeHTTP(getRec, getReq)
	if getRec.Code != http.StatusOK {
		t.Fatalf("expected sqlite-backed artifact get 200, got %d body=%s", getRec.Code, getRec.Body.String())
	}

	listReq := httptest.NewRequest(http.MethodGet, "/v1/artifacts?job_id="+envelope.Data.JobID, nil)
	listReq.Header.Set("Authorization", "Bearer "+sessionToken)
	listRec := httptest.NewRecorder()
	server.Handler().ServeHTTP(listRec, listReq)
	if listRec.Code != http.StatusOK {
		t.Fatalf("expected sqlite-backed artifacts list 200, got %d body=%s", listRec.Code, listRec.Body.String())
	}
}

func TestResolveImageToImageReferenceAcceptsCanonicalImageURL(t *testing.T) {
	cfg := Config{
		SQLitePath:    filepath.Join(t.TempDir(), "popiart.db"),
		SkillhubDir:   makeEmptySkillhub(t),
		SessionSecret: "test-secret",
	}
	server, err := NewWithConfig(cfg)
	if err != nil {
		t.Fatalf("NewWithConfig: %v", err)
	}

	imageBytes := tinyPNG(t)
	refSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "image/png")
		_, _ = w.Write(imageBytes)
	}))
	defer refSrv.Close()

	editRef, err := server.resolveImageToImageReference(context.Background(), &job{
		UserID:      "user_test",
		UpstreamKey: "sk-upstream",
		SkillID:     "popiskill-image-img2img-basic-v1",
	}, map[string]any{
		"image": refSrv.URL + "/reference.png",
	})
	if err != nil {
		t.Fatalf("resolveImageToImageReference: %v", err)
	}
	if editRef.URL != refSrv.URL+"/reference.png" {
		t.Fatalf("expected resolved URL to match input image, got %q", editRef.URL)
	}
	if editRef.ContentType != "image/png" {
		t.Fatalf("expected image/png, got %q", editRef.ContentType)
	}
	if !bytes.Equal(editRef.Content, imageBytes) {
		t.Fatal("expected resolved canonical image content to match downloaded bytes")
	}
}

func TestResolveImageToImageReferenceAcceptsSignedSameOriginMediaURL(t *testing.T) {
	cfg := Config{
		SQLitePath:    filepath.Join(t.TempDir(), "popiart.db"),
		SkillhubDir:   makeEmptySkillhub(t),
		SessionSecret: "test-secret",
	}
	server, err := NewWithConfig(cfg)
	if err != nil {
		t.Fatalf("NewWithConfig: %v", err)
	}
	srv := httptest.NewServer(server.Handler())
	defer srv.Close()
	server.cfg.PublicBaseURL = srv.URL

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

	imageBytes := tinyPNG(t)
	var body bytes.Buffer
	writer := multipart.NewWriter(&body)
	part, err := writer.CreateFormFile("file", "same-origin.png")
	if err != nil {
		t.Fatalf("create form file: %v", err)
	}
	if _, err := part.Write(imageBytes); err != nil {
		t.Fatalf("write form file: %v", err)
	}
	if err := writer.Close(); err != nil {
		t.Fatalf("close multipart writer: %v", err)
	}

	req := httptest.NewRequest(http.MethodPost, "/v1/media/upload", &body)
	req.Header.Set("Authorization", "Bearer "+sessionToken)
	req.Header.Set("Content-Type", writer.FormDataContentType())
	rec := httptest.NewRecorder()
	server.Handler().ServeHTTP(rec, req)
	if rec.Code != http.StatusCreated {
		t.Fatalf("expected media upload 201, got %d body=%s", rec.Code, rec.Body.String())
	}

	var envelope struct {
		OK   bool  `json:"ok"`
		Data media `json:"data"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &envelope); err != nil {
		t.Fatalf("decode media upload response: %v body=%s", err, rec.Body.String())
	}
	if !envelope.OK || envelope.Data.URL == "" {
		t.Fatalf("expected signed media url, got %s", rec.Body.String())
	}

	editRef, err := server.resolveImageToImageReference(context.Background(), &job{
		UserID:      current.User.ID,
		SessionID:   current.Token,
		UpstreamKey: current.UpstreamKey,
		SkillID:     "popiskill-image-img2img-basic-v1",
	}, map[string]any{
		"image": envelope.Data.URL,
	})
	if err != nil {
		t.Fatalf("resolveImageToImageReference same-origin: %v", err)
	}
	if editRef.URL != envelope.Data.URL {
		t.Fatalf("expected resolved URL to round-trip, got %q", editRef.URL)
	}
	if editRef.ContentType != "image/png" {
		t.Fatalf("expected image/png, got %q", editRef.ContentType)
	}
	if !bytes.Equal(editRef.Content, imageBytes) {
		t.Fatal("expected same-origin signed media bytes to match uploaded file")
	}
}