250 lines
8.5 KiB
Go
250 lines
8.5 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"reflect"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func testPersistedMesh(textured bool) *meshData {
|
|
m := &meshData{
|
|
NVerts: 3, NTris: 1,
|
|
Verts: []float32{0, 0, 0, 1, 0, 0, 0, 1, 0},
|
|
Normals: []float32{0, 0, 1, 0, 0, 1, 0, 0, 1},
|
|
Tris: []int32{0, 1, 2},
|
|
}
|
|
if textured {
|
|
m.PBR = []float32{
|
|
1, 0, 0, 0.1, 0.2, 1,
|
|
0, 1, 0, 0.3, 0.4, 0.9,
|
|
0, 0, 1, 0.5, 0.6, 0.8,
|
|
}
|
|
}
|
|
return m
|
|
}
|
|
|
|
func TestMeshFileRoundTrip(t *testing.T) {
|
|
for _, textured := range []bool{false, true} {
|
|
t.Run(map[bool]string{false: "geometry", true: "pbr"}[textured], func(t *testing.T) {
|
|
path := filepath.Join(t.TempDir(), "mesh.t2mesh")
|
|
want := testPersistedMesh(textured)
|
|
if err := writeMeshFile(path, want); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
got, err := readMeshFile(path)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !reflect.DeepEqual(got, want) {
|
|
t.Fatalf("mesh round trip mismatch\n got: %#v\nwant: %#v", got, want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestPersistAndRestoreCompletedJob(t *testing.T) {
|
|
store := t.TempDir()
|
|
s := &server{jobs: map[string]*job{}, storeDir: store}
|
|
wantFrames := []frameMeta{
|
|
{Stage: "sampling sparse structure", Step: 1, Total: 2, Kind: "voxel"},
|
|
{Stage: "sampling shape SLAT", Step: 2, Total: 2, Kind: "mesh"},
|
|
}
|
|
wantPreviews := [][]byte{[]byte("T2VOX01-frame"), []byte("T2MESH01-frame")}
|
|
j := &job{
|
|
ID: "0123456789abcdef", State: "running", CreatedAt: 123456789,
|
|
StartedAt: 123456800, FinishedAt: 123499000, DurationMS: 42200,
|
|
Quality: "1024", Thumbnail: "data:image/jpeg;base64,dGVzdA==",
|
|
PreviewSeq: len(wantFrames), Frames: wantFrames, LivePreview: true,
|
|
StageTimings: []stageTiming{{Stage: "sampling sparse structure", Milliseconds: 21000}},
|
|
pipeline: pipe1024, seed: 42, steps: 12, textureSteps: 10, guidance: 7.5,
|
|
previews: wantPreviews, mesh: testPersistedMesh(true),
|
|
image: []byte("processed-input"), source: []byte("exact-original"),
|
|
}
|
|
if err := s.persistJob(j); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(j.previews) != 0 || j.persistDir == "" || j.meshPath == "" ||
|
|
j.inputPath == "" || j.sourcePath == "" {
|
|
t.Fatalf("persist did not switch assets to disk: previews=%d dir=%q mesh=%q input=%q source=%q",
|
|
len(j.previews), j.persistDir, j.meshPath, j.inputPath, j.sourcePath)
|
|
}
|
|
|
|
restarted := &server{jobs: map[string]*job{}, storeDir: store}
|
|
n, err := restarted.restoreJobs()
|
|
if err != nil || n != 1 {
|
|
t.Fatalf("restoreJobs() = %d, %v; want 1, nil", n, err)
|
|
}
|
|
got := restarted.jobs[j.ID]
|
|
if got == nil || got.State != "done" || got.Quality != "1024" ||
|
|
got.Thumbnail != j.Thumbnail || got.CreatedAt != j.CreatedAt ||
|
|
got.StartedAt != j.StartedAt || got.FinishedAt != j.FinishedAt ||
|
|
got.DurationMS != j.DurationMS || got.LivePreview != j.LivePreview ||
|
|
!reflect.DeepEqual(got.StageTimings, j.StageTimings) {
|
|
t.Fatalf("restored metadata = %#v", got)
|
|
}
|
|
if got.mesh != nil || len(got.previews) != 0 {
|
|
t.Fatal("restored binary assets should remain lazy until requested")
|
|
}
|
|
if source, err := os.ReadFile(got.sourcePath); err != nil || string(source) != "exact-original" {
|
|
t.Fatalf("restored source image = %q, %v", source, err)
|
|
}
|
|
if input, err := os.ReadFile(got.inputPath); err != nil || string(input) != "processed-input" {
|
|
t.Fatalf("restored generation input = %q, %v", input, err)
|
|
}
|
|
rr := httptest.NewRecorder()
|
|
restarted.handleSource(rr, httptest.NewRequest(http.MethodGet, "/api/source/"+j.ID, nil))
|
|
if rr.Code != http.StatusOK || rr.Body.String() != "exact-original" {
|
|
t.Fatalf("GET source = %d %q", rr.Code, rr.Body.String())
|
|
}
|
|
mesh, err := restarted.loadJobMesh(got)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !reflect.DeepEqual(mesh, j.mesh) {
|
|
t.Fatal("restored mesh mismatch")
|
|
}
|
|
for i, want := range wantPreviews {
|
|
blob, err := loadJobPreview(got, i)
|
|
if err != nil {
|
|
t.Fatalf("preview %d: %v", i, err)
|
|
}
|
|
if !bytes.Equal(blob, want) {
|
|
t.Fatalf("preview %d = %q, want %q", i, blob, want)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestRegenerateUsesPersistedInputWithoutUpload(t *testing.T) {
|
|
dir := t.TempDir()
|
|
inputPath := filepath.Join(dir, "input.img")
|
|
sourcePath := filepath.Join(dir, "source.img")
|
|
if err := os.WriteFile(inputPath, []byte("processed-input"), 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := os.WriteFile(sourcePath, []byte("exact-original"), 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
old := &job{
|
|
ID: "old", State: "done", CreatedAt: 1, Quality: "512", Thumbnail: "thumb",
|
|
inputPath: inputPath, sourcePath: sourcePath, pipeline: pipe512,
|
|
seed: 7, steps: 8, textureSteps: 9, guidance: 6.5,
|
|
}
|
|
s := &server{jobs: map[string]*job{"old": old}, q: make(chan *job, 1)}
|
|
req := httptest.NewRequest(http.MethodPost, "/api/regenerate/old",
|
|
strings.NewReader("quality=1024&seed=42&steps=14&texture_steps=15&guidance=8&preview=0"))
|
|
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
|
rr := httptest.NewRecorder()
|
|
s.handleRegenerate(rr, req)
|
|
if rr.Code != http.StatusOK {
|
|
t.Fatalf("POST regenerate = %d: %s", rr.Code, rr.Body.String())
|
|
}
|
|
var response map[string]string
|
|
if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
regenerated := s.jobs[response["job"]]
|
|
if regenerated == nil {
|
|
t.Fatal("regenerated job was not queued")
|
|
}
|
|
if string(regenerated.image) != "processed-input" || string(regenerated.source) != "exact-original" {
|
|
t.Fatalf("regenerated bytes = input %q, source %q", regenerated.image, regenerated.source)
|
|
}
|
|
if regenerated.Quality != "1024" || regenerated.pipeline != pipe1024 ||
|
|
regenerated.background != backgroundKeep || regenerated.seed != 42 ||
|
|
regenerated.steps != 14 || regenerated.textureSteps != 15 ||
|
|
regenerated.guidance != 8 || regenerated.LivePreview {
|
|
t.Fatalf("regenerated settings = %#v", regenerated)
|
|
}
|
|
}
|
|
|
|
func TestGenerateKeepsOriginalSeparateFromProcessedInput(t *testing.T) {
|
|
var body bytes.Buffer
|
|
mw := multipart.NewWriter(&body)
|
|
inputPart, err := mw.CreateFormFile("image", "input.png")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
inputPart.Write([]byte("processed-png"))
|
|
sourcePart, err := mw.CreateFormFile("source", "camera-original.webp")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
sourcePart.Write([]byte("exact-original-webp"))
|
|
if err := mw.Close(); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
s := &server{jobs: map[string]*job{}, q: make(chan *job, 1)}
|
|
req := httptest.NewRequest(http.MethodPost, "/api/generate", &body)
|
|
req.Header.Set("Content-Type", mw.FormDataContentType())
|
|
rr := httptest.NewRecorder()
|
|
s.handleGenerate(rr, req)
|
|
if rr.Code != http.StatusOK {
|
|
t.Fatalf("POST generate = %d: %s", rr.Code, rr.Body.String())
|
|
}
|
|
var response map[string]string
|
|
if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
j := s.jobs[response["job"]]
|
|
if j == nil || string(j.image) != "processed-png" || string(j.source) != "exact-original-webp" {
|
|
t.Fatalf("queued job bytes = %#v", j)
|
|
}
|
|
}
|
|
|
|
func TestServerHistoryListAndDelete(t *testing.T) {
|
|
root := t.TempDir()
|
|
oldDir := filepath.Join(root, "old")
|
|
if err := os.Mkdir(oldDir, 0o755); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
old := &job{ID: "old", State: "done", CreatedAt: 10, Quality: "512", persistDir: oldDir}
|
|
newer := &job{ID: "new", State: "done", CreatedAt: 20, Quality: "1024", Thumbnail: "thumb"}
|
|
active := &job{ID: "active", State: "running", CreatedAt: 30}
|
|
s := &server{jobs: map[string]*job{"old": old, "new": newer, "active": active}}
|
|
|
|
rr := httptest.NewRecorder()
|
|
s.handleJobs(rr, httptest.NewRequest(http.MethodGet, "/api/jobs", nil))
|
|
if rr.Code != http.StatusOK {
|
|
t.Fatalf("GET /api/jobs = %d: %s", rr.Code, rr.Body.String())
|
|
}
|
|
var got []jobSummary
|
|
if err := json.Unmarshal(rr.Body.Bytes(), &got); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(got) != 2 || got[0].ID != "new" || got[1].ID != "old" {
|
|
t.Fatalf("history order/content = %#v", got)
|
|
}
|
|
|
|
rr = httptest.NewRecorder()
|
|
s.handleJob(rr, httptest.NewRequest(http.MethodDelete, "/api/job/old", nil))
|
|
if rr.Code != http.StatusNoContent {
|
|
t.Fatalf("DELETE /api/job/old = %d: %s", rr.Code, rr.Body.String())
|
|
}
|
|
if s.jobs["old"] != nil {
|
|
t.Fatal("deleted job remains in server index")
|
|
}
|
|
if _, err := os.Stat(oldDir); !os.IsNotExist(err) {
|
|
t.Fatalf("deleted job directory still exists: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestRestoreIgnoresIncompleteTemporaryDirectory(t *testing.T) {
|
|
store := t.TempDir()
|
|
if err := os.Mkdir(filepath.Join(store, ".unfinished.tmp-123"), 0o755); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
s := &server{jobs: map[string]*job{}, storeDir: store}
|
|
n, err := s.restoreJobs()
|
|
if err != nil || n != 0 {
|
|
t.Fatalf("restoreJobs() = %d, %v; want 0, nil", n, err)
|
|
}
|
|
}
|