easyai-ai-gateway/apps/api/internal/acceptanceemulator/server_test.go

package acceptanceemulator

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"image"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"
)

func TestGeminiAndVolcesProtocolEmulation(t *testing.T) {
	var mu sync.Mutex
	now := time.Unix(1_800_000_000, 0)
	server := New(Config{
		Now: func() time.Time {
			mu.Lock()
			defer mu.Unlock()
			return now
		},
		Wait: func(context.Context, time.Duration) error { return nil },
	})
	httpServer := httptest.NewServer(server.Handler())
	defer httpServer.Close()

	input := paddedPNG(256 << 10)
	geminiBody, _ := json.Marshal(map[string]any{
		"contents": []any{map[string]any{"parts": []any{
			map[string]any{"text": "edit image"},
			map[string]any{"inlineData": map[string]any{
				"mimeType": "image/png", "data": base64.StdEncoding.EncodeToString(input),
			}},
		}}},
		"generationConfig": map[string]any{"responseModalities": []any{"IMAGE"}},
	})
	response, err := http.Post(httpServer.URL+"/v1beta/models/gemini-test:generateContent", "application/json", bytes.NewReader(geminiBody))
	if err != nil {
		t.Fatalf("Gemini request: %v", err)
	}
	if response.StatusCode != http.StatusOK {
		payload, _ := io.ReadAll(response.Body)
		t.Fatalf("Gemini status=%d body=%s", response.StatusCode, payload)
	}
	var geminiResult map[string]any
	if err := json.NewDecoder(response.Body).Decode(&geminiResult); err != nil {
		t.Fatalf("decode Gemini response: %v", err)
	}
	_ = response.Body.Close()

	fixtureResponse, err := http.Get(httpServer.URL + "/fixtures/image-12-4k.jpg")
	if err != nil {
		t.Fatalf("get 4K fixture: %v", err)
	}
	config, format, err := image.DecodeConfig(fixtureResponse.Body)
	_ = fixtureResponse.Body.Close()
	if err != nil || format != "jpeg" || config.Width != 4096 || config.Height != 2160 {
		t.Fatalf("4K fixture format=%s config=%+v err=%v", format, config, err)
	}

	content := []any{map[string]any{"type": "text", "text": "video"}}
	for index, name := range []string{"image-00.png", "image-04.jpg", "image-08.webp"} {
		role := "reference_image"
		if index == 0 {
			role = "first_frame"
		}
		content = append(content, map[string]any{
			"type": "image_url", "role": role,
			"image_url": map[string]any{"url": httpServer.URL + "/fixtures/" + name},
		})
	}
	videoBody, _ := json.Marshal(map[string]any{"model": "seedance-test", "content": content, "seed": 1})
	response, err = http.Post(httpServer.URL+"/contents/generations/tasks", "application/json", bytes.NewReader(videoBody))
	if err != nil {
		t.Fatalf("submit video: %v", err)
	}
	if response.StatusCode != http.StatusOK {
		t.Fatalf("submit video status=%d", response.StatusCode)
	}
	var submitted map[string]any
	_ = json.NewDecoder(response.Body).Decode(&submitted)
	_ = response.Body.Close()
	taskID := strings.TrimSpace(stringValue(submitted["id"]))
	if taskID == "" {
		t.Fatal("video emulator returned no task ID")
	}
	mu.Lock()
	now = now.Add(20 * time.Second)
	mu.Unlock()
	response, err = http.Get(httpServer.URL + "/contents/generations/tasks/" + taskID)
	if err != nil {
		t.Fatalf("poll video: %v", err)
	}
	var polled map[string]any
	_ = json.NewDecoder(response.Body).Decode(&polled)
	_ = response.Body.Close()
	if polled["status"] != "succeeded" {
		t.Fatalf("video status=%v", polled["status"])
	}

	response, err = http.Get(httpServer.URL + "/report")
	if err != nil {
		t.Fatalf("get report: %v", err)
	}
	var report Report
	_ = json.NewDecoder(response.Body).Decode(&report)
	_ = response.Body.Close()
	if report.GeminiRequests != 1 || report.VideoSubmissions != 1 || report.VideoInvalid != 0 ||
		report.VideoReferenceCounts["3"] != 1 || report.UniqueImageHashes != 3 {
		t.Fatalf("unexpected emulator report: %+v", report)
	}
}

func TestPaddedPNGIsValidAndExact(t *testing.T) {
	for _, size := range []int{256 << 10, 4 << 20, 8 << 20} {
		payload := paddedPNG(size)
		if len(payload) != size {
			t.Fatalf("PNG bytes=%d, want=%d", len(payload), size)
		}
		if _, _, err := image.DecodeConfig(bytes.NewReader(payload)); err != nil {
			t.Fatalf("decode padded PNG: %v", err)
		}
	}
}

func TestScaledDelayStaysOnWholeSecondsWithinProfile(t *testing.T) {
	delay := scaledDelay(8*time.Second, 15*time.Second, 2<<20)
	if delay < 8*time.Second || delay > 15*time.Second || delay%time.Second != 0 {
		t.Fatalf("scaled delay=%s", delay)
	}
}

func TestVolcesProtocolAcceptsThreeSixAndNineReferenceImages(t *testing.T) {
	server := New(Config{Wait: func(context.Context, time.Duration) error { return nil }})
	httpServer := httptest.NewServer(server.Handler())
	defer httpServer.Close()

	for _, imageCount := range []int{3, 6, 9} {
		content := []any{map[string]any{"type": "text", "text": "multi-reference video"}}
		for index := 0; index < imageCount; index++ {
			role := "reference_image"
			if index == 0 {
				role = "first_frame"
			}
			if imageCount > 3 && index == imageCount-1 {
				role = "last_frame"
			}
			content = append(content, map[string]any{
				"type": "image_url", "role": role,
				"image_url": map[string]any{
					"url": fmt.Sprintf("%s/fixtures/image-%02d.png", httpServer.URL, index%4),
				},
			})
		}
		body, _ := json.Marshal(map[string]any{"model": "seedance-test", "content": content, "seed": imageCount})
		response, err := http.Post(httpServer.URL+"/contents/generations/tasks", "application/json", bytes.NewReader(body))
		if err != nil {
			t.Fatalf("%d-image submit: %v", imageCount, err)
		}
		_ = response.Body.Close()
		if response.StatusCode != http.StatusOK {
			t.Fatalf("%d-image status=%d", imageCount, response.StatusCode)
		}
	}

	response, err := http.Get(httpServer.URL + "/report")
	if err != nil {
		t.Fatalf("get report: %v", err)
	}
	defer response.Body.Close()
	var report Report
	if err := json.NewDecoder(response.Body).Decode(&report); err != nil {
		t.Fatalf("decode report: %v", err)
	}
	for _, imageCount := range []string{"3", "6", "9"} {
		if report.VideoReferenceCounts[imageCount] != 1 {
			t.Fatalf("reference counts=%v", report.VideoReferenceCounts)
		}
	}
}