mAi
64120c27d7
Schema (applied via migration imagen_series_init): - imagen.series parent table (prompt + params + count CHECK 1..10 + selected_image_id) - imagen.jobs += series_id (FK) + series_idx - imagen.images += series_id (FK) - Owner-scoped RLS on series (SELECT/INSERT/UPDATE) + grants - Partial indexes WHERE series_id IS NOT NULL on both child tables Worker pipeline: - worker.Job += SeriesID, populated from imagen.jobs.series_id via the claim query. - cloud.SyncRequest += SeriesID; insertRow writes series_id when non-empty, omits the key when empty so solo runs leave the column NULL. - maybeCloudSync threads seriesID from job.SeriesID through to the cloud sink. generate.go (CLI) always passes "" — solo path unchanged. Tests: - worker: SeriesID propagates from Job to fakePipeline.lastJob unchanged, solo job keeps it empty. - cloud: SyncRequest.SeriesID lands as row.series_id in the POST body; empty SeriesID omits the key entirely. Refs ImaGen#9.
375 lines
11 KiB
Go
375 lines
11 KiB
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)
|
|
}
|
|
}
|
|
|
|
// TestSyncWritesSeriesID is the second half of the ImaGen#9 propagation
|
|
// contract: when SeriesID is non-empty, the POST body to imagen.images
|
|
// carries `series_id`. When empty, the key is omitted entirely so the
|
|
// row's series_id stays NULL (solo-run path, list-page query
|
|
// `WHERE series_id IS NULL` keeps showing it).
|
|
func TestSyncWritesSeriesID(t *testing.T) {
|
|
const seriesID = "22222222-2222-2222-2222-222222222222"
|
|
f := newFakeSupabase(t)
|
|
s := newSink(f.server)
|
|
|
|
_, err := s.Sync(context.Background(), SyncRequest{
|
|
Date: "2026-05-11", Slug: "x", Seed: 1, Ext: "png",
|
|
PNG: []byte("p"), Prompt: "p", Backend: "b",
|
|
SeriesID: seriesID,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Sync: %v", err)
|
|
}
|
|
var row map[string]any
|
|
if err := json.Unmarshal(f.insertBody, &row); err != nil {
|
|
t.Fatalf("parse insert body: %v\n%s", err, f.insertBody)
|
|
}
|
|
if row["series_id"] != seriesID {
|
|
t.Fatalf("row.series_id = %v want %q", row["series_id"], seriesID)
|
|
}
|
|
}
|
|
|
|
func TestSyncOmitsSeriesIDWhenEmpty(t *testing.T) {
|
|
f := newFakeSupabase(t)
|
|
s := newSink(f.server)
|
|
|
|
_, err := s.Sync(context.Background(), SyncRequest{
|
|
Date: "2026-05-11", Slug: "x", Seed: 1, Ext: "png",
|
|
PNG: []byte("p"), Prompt: "p", Backend: "b",
|
|
// SeriesID intentionally empty.
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Sync: %v", err)
|
|
}
|
|
var row map[string]any
|
|
if err := json.Unmarshal(f.insertBody, &row); err != nil {
|
|
t.Fatalf("parse insert body: %v\n%s", err, f.insertBody)
|
|
}
|
|
if _, present := row["series_id"]; present {
|
|
t.Fatalf("solo run should omit series_id from POST body, got %v", row["series_id"])
|
|
}
|
|
}
|
|
|
|
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)
|
|
}
|
|
}
|
|
}
|
|
}
|