ImaGen/internal/cloud/cloud_test.go

package cloud

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

// fakeSupabase is a tiny stand-in for Supabase Storage + PostgREST. It
// records what came in and returns canned responses based on path.
type fakeSupabase struct {
	t           *testing.T
	mux         *http.ServeMux
	server      *httptest.Server
	uploadCalls int32
	insertCalls int32
	uploadBytes []byte
	uploadHdr   http.Header
	insertBody  []byte
	insertHdr   http.Header
}

func newFakeSupabase(t *testing.T, opts ...func(*fakeSupabase)) *fakeSupabase {
	f := &fakeSupabase{t: t}
	f.mux = http.NewServeMux()
	// Storage upload — anything under /storage/v1/object/<bucket>/...
	f.mux.HandleFunc("/storage/v1/object/imagen-generated/", func(w http.ResponseWriter, r *http.Request) {
		atomic.AddInt32(&f.uploadCalls, 1)
		body, _ := io.ReadAll(r.Body)
		f.uploadBytes = body
		f.uploadHdr = r.Header.Clone()
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(`{"Key":"imagen-generated/somepath"}`))
	})
	// Storage sign URL
	f.mux.HandleFunc("/storage/v1/object/sign/imagen-generated/", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.Write([]byte(`{"signedURL":"/storage/v1/object/sign/imagen-generated/some.png?token=abc"}`))
	})
	// PostgREST insert
	f.mux.HandleFunc("/rest/v1/images", func(w http.ResponseWriter, r *http.Request) {
		atomic.AddInt32(&f.insertCalls, 1)
		body, _ := io.ReadAll(r.Body)
		f.insertBody = body
		f.insertHdr = r.Header.Clone()
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusCreated)
		w.Write([]byte(`[{"id":"00000000-0000-0000-0000-000000000abc"}]`))
	})
	for _, opt := range opts {
		opt(f)
	}
	f.server = httptest.NewServer(f.mux)
	t.Cleanup(f.server.Close)
	return f
}

func newSink(server *httptest.Server) *Sink {
	return &Sink{
		URL:            server.URL,
		APIKey:         "fake-service-key",
		OwnerUserID:    "00000000-0000-0000-0000-000000000001",
		HTTP:           server.Client(),
		MaxRetries:     2,
		InitialBackoff: time.Millisecond,
	}
}

func TestSyncHappyPath(t *testing.T) {
	f := newFakeSupabase(t)
	s := newSink(f.server)

	cost := 0.003
	res, err := s.Sync(context.Background(), SyncRequest{
		Date:            "2026-05-11",
		Slug:            "lighthouse",
		Seed:            42,
		Ext:             "png",
		PNG:             []byte("PNGbytes"),
		MimeType:        "image/png",
		Prompt:          "a tiny lighthouse on a stormy cliff",
		Backend:         "flux-schnell-local",
		Model:           "flux1-schnell",
		Steps:           4,
		Width:           1024,
		Height:          1024,
		LatencyMs:       1500,
		CostUSDEstimate: &cost,
		Sidecar: map[string]any{
			"timestamp": "2026-05-11T01:30:00Z",
			"backend":   "flux-schnell-local",
		},
	})
	if err != nil {
		t.Fatalf("Sync: %v", err)
	}
	if res.StoragePath != "2026-05-11/lighthouse-42.png" {
		t.Errorf("storage_path = %q", res.StoragePath)
	}
	if res.ImageID != "00000000-0000-0000-0000-000000000abc" {
		t.Errorf("image_id = %q", res.ImageID)
	}
	if got := atomic.LoadInt32(&f.uploadCalls); got != 1 {
		t.Errorf("upload calls = %d, want 1", got)
	}
	if got := atomic.LoadInt32(&f.insertCalls); got != 1 {
		t.Errorf("insert calls = %d, want 1", got)
	}
	if !bytes.Equal(f.uploadBytes, []byte("PNGbytes")) {
		t.Errorf("uploaded bytes = %q", f.uploadBytes)
	}

	// Verify the row payload carries the prompt + computed hash + non-zero
	// metadata. Empty fields should be omitted from the JSON body so RLS
	// won't see surprise keys.
	var row map[string]any
	if err := json.Unmarshal(f.insertBody, &row); err != nil {
		t.Fatalf("insert body parse: %v\n%s", err, f.insertBody)
	}
	if row["prompt"] != "a tiny lighthouse on a stormy cliff" {
		t.Errorf("row.prompt = %v", row["prompt"])
	}
	if row["owner_user_id"] != "00000000-0000-0000-0000-000000000001" {
		t.Errorf("row.owner_user_id = %v", row["owner_user_id"])
	}
	if row["storage_path"] != "2026-05-11/lighthouse-42.png" {
		t.Errorf("row.storage_path = %v", row["storage_path"])
	}
	hash, _ := row["prompt_hash"].(string)
	if len(hash) != 64 {
		t.Errorf("prompt_hash should be 64-char sha256 hex, got %q", hash)
	}
	if row["backend"] != "flux-schnell-local" {
		t.Errorf("row.backend = %v", row["backend"])
	}
	if row["seed"].(float64) != 42 {
		t.Errorf("row.seed = %v", row["seed"])
	}
	if row["latency_ms"].(float64) != 1500 {
		t.Errorf("row.latency_ms = %v", row["latency_ms"])
	}
	if row["cost_usd_estimate"].(float64) != 0.003 {
		t.Errorf("row.cost = %v", row["cost_usd_estimate"])
	}
	if row["sidecar"] == nil {
		t.Errorf("row.sidecar missing")
	}

	// PostgREST schema headers — hardcoded to "imagen".
	if got := f.insertHdr.Get("Accept-Profile"); got != "imagen" {
		t.Errorf("Accept-Profile = %q", got)
	}
	if got := f.insertHdr.Get("Content-Profile"); got != "imagen" {
		t.Errorf("Content-Profile = %q", got)
	}
	if got := f.insertHdr.Get("Authorization"); !strings.HasPrefix(got, "Bearer ") {
		t.Errorf("Authorization = %q", got)
	}

	// Storage upsert should be set so re-runs of the same date+slug+seed
	// don't fail with 409.
	if got := f.uploadHdr.Get("x-upsert"); got != "true" {
		t.Errorf("x-upsert = %q", got)
	}
}

func TestSyncRetryOn5xx(t *testing.T) {
	var uploadAttempts int32
	mux := http.NewServeMux()
	mux.HandleFunc("/storage/v1/object/imagen-generated/", func(w http.ResponseWriter, r *http.Request) {
		n := atomic.AddInt32(&uploadAttempts, 1)
		// Two 503s, then OK.
		if n < 3 {
			http.Error(w, "service unavailable", http.StatusServiceUnavailable)
			return
		}
		w.WriteHeader(http.StatusOK)
	})
	mux.HandleFunc("/rest/v1/images", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusCreated)
		w.Write([]byte(`[{"id":"row-id"}]`))
	})
	srv := httptest.NewServer(mux)
	defer srv.Close()
	s := newSink(srv)

	res, err := s.Sync(context.Background(), SyncRequest{
		Date: "2026-05-11", Slug: "x", Seed: 1, Ext: "png",
		PNG: []byte("p"), Prompt: "p", Backend: "b",
	})
	if err != nil {
		t.Fatalf("Sync (with retry): %v", err)
	}
	if got := atomic.LoadInt32(&uploadAttempts); got != 3 {
		t.Errorf("upload attempts = %d, want 3", got)
	}
	if res.ImageID != "row-id" {
		t.Errorf("image_id = %q", res.ImageID)
	}
}

func TestSyncNoRetryOn4xx(t *testing.T) {
	var uploadAttempts int32
	mux := http.NewServeMux()
	mux.HandleFunc("/storage/v1/object/imagen-generated/", func(w http.ResponseWriter, r *http.Request) {
		atomic.AddInt32(&uploadAttempts, 1)
		http.Error(w, `{"message":"bad request"}`, http.StatusBadRequest)
	})
	srv := httptest.NewServer(mux)
	defer srv.Close()
	s := newSink(srv)

	_, err := s.Sync(context.Background(), SyncRequest{
		Date: "2026-05-11", Slug: "x", Seed: 1, Ext: "png",
		PNG: []byte("p"), Prompt: "p", Backend: "b",
	})
	if err == nil {
		t.Fatal("expected error on 400")
	}
	if !strings.Contains(err.Error(), "400") {
		t.Errorf("error should mention 400 status: %v", err)
	}
	if got := atomic.LoadInt32(&uploadAttempts); got != 1 {
		t.Errorf("upload attempts = %d, want 1 (no retry on 4xx)", got)
	}
}

func TestSyncMissingOwnerUserID(t *testing.T) {
	srv := httptest.NewServer(http.NewServeMux())
	defer srv.Close()
	s := &Sink{
		URL:    srv.URL,
		APIKey: "k",
		// OwnerUserID intentionally empty.
		HTTP:           srv.Client(),
		InitialBackoff: time.Millisecond,
	}
	_, err := s.Sync(context.Background(), SyncRequest{
		Date: "2026-05-11", Slug: "x", Seed: 1, Ext: "png",
		PNG: []byte("p"), Prompt: "p", Backend: "b",
	})
	if err == nil {
		t.Fatal("expected error when owner_user_id unset")
	}
	if !strings.Contains(err.Error(), "owner_user_id") {
		t.Errorf("error should mention owner_user_id: %v", err)
	}
}

func TestSyncRequiresDateAndSlug(t *testing.T) {
	srv := httptest.NewServer(http.NewServeMux())
	defer srv.Close()
	s := newSink(srv)
	_, err := s.Sync(context.Background(), SyncRequest{
		Slug: "x", Seed: 1, Ext: "png",
		PNG: []byte("p"), Prompt: "p", Backend: "b",
	})
	if err == nil {
		t.Fatal("expected error for missing date")
	}
}

func TestSignedURL(t *testing.T) {
	f := newFakeSupabase(t)
	s := newSink(f.server)
	got, err := s.SignedURL(context.Background(), "2026-05-11/x.png", 60)
	if err != nil {
		t.Fatalf("SignedURL: %v", err)
	}
	want := f.server.URL + "/storage/v1/object/sign/imagen-generated/some.png?token=abc"
	if got != want {
		t.Errorf("signed URL = %q, want %q", got, want)
	}
}

func TestSyncDBFailureSurfacesPathOnError(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("/storage/v1/object/imagen-generated/", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	})
	mux.HandleFunc("/rest/v1/images", func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "schema cache miss", http.StatusInternalServerError)
	})
	srv := httptest.NewServer(mux)
	defer srv.Close()
	s := newSink(srv)
	res, err := s.Sync(context.Background(), SyncRequest{
		Date: "2026-05-11", Slug: "x", Seed: 9, Ext: "png",
		PNG: []byte("p"), Prompt: "p", Backend: "b",
	})
	if err == nil {
		t.Fatal("expected error from DB insert failure")
	}
	// Storage upload succeeded — caller can still see the upload landed.
	if res == nil || res.StoragePath != "2026-05-11/x-9.png" {
		t.Errorf("expected storage_path on partial success, got %+v", res)
	}
}

func TestPathEscape(t *testing.T) {
	cases := map[string]string{
		"2026-05-11/lighthouse-42.png":              "2026-05-11/lighthouse-42.png",
		"2026-05-11/two words.png":                  "2026-05-11/two%20words.png",
		"with#hash/and?query.png":                   "with%23hash/and%3Fquery.png",
	}
	for in, want := range cases {
		got := pathEscape(in)
		if got != want {
			t.Errorf("pathEscape(%q) = %q, want %q", in, got, want)
		}
		// Sanity: every part should round-trip via url.PathUnescape.
		for _, seg := range strings.Split(got, "/") {
			if _, err := url.PathUnescape(seg); err != nil {
				t.Errorf("segment %q failed unescape: %v", seg, err)
			}
		}
	}
}