ImaGen/internal/backend/mock.go

package backend

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/binary"
	"fmt"
	"image"
	"image/color"
	"image/png"
	"io"
	"math/rand"
	"time"
)

// Mock is a deterministic image generator used for tests, smoke checks and as
// the reference implementation of the Backend contract. Same seed + same prompt
// always yields the same PNG. It does not call the network.
type Mock struct {
	instance string
}

// NewMock builds a Mock backend. cfg is accepted for symmetry with real
// adapters but is ignored.
func NewMock(name string, _ map[string]any) (Backend, error) {
	if name == "" {
		name = "mock"
	}
	return &Mock{instance: name}, nil
}

// Name returns the user-facing instance name.
func (m *Mock) Name() string { return m.instance }

// Generate paints a deterministic gradient sized to req.Width×req.Height and
// returns it as a PNG. Width/Height default to 256 when zero.
func (m *Mock) Generate(ctx context.Context, req Request) (*Result, error) {
	w, h := req.Width, req.Height
	if w == 0 {
		w = 256
	}
	if h == 0 {
		h = 256
	}
	if w < 1 || h < 1 || w > 8192 || h > 8192 {
		return nil, fmt.Errorf("mock: invalid size %dx%d", w, h)
	}

	seed := req.Seed
	if seed == 0 {
		seed = derivedSeed(req.Prompt)
	}
	rng := rand.New(rand.NewSource(seed))
	baseR := uint8(rng.Intn(256))
	baseG := uint8(rng.Intn(256))
	baseB := uint8(rng.Intn(256))

	start := time.Now()
	img := image.NewRGBA(image.Rect(0, 0, w, h))
	for y := 0; y < h; y++ {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
		}
		for x := 0; x < w; x++ {
			fx := float64(x) / float64(w)
			fy := float64(y) / float64(h)
			img.Set(x, y, color.RGBA{
				R: blend(baseR, fx),
				G: blend(baseG, fy),
				B: blend(baseB, (fx+fy)/2),
				A: 255,
			})
		}
	}
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		return nil, fmt.Errorf("mock: encode png: %w", err)
	}
	return &Result{
		ImageReader: io.NopCloser(&buf),
		MimeType:    "image/png",
		Metadata: map[string]any{
			"backend":     m.instance,
			"backend_type": "mock",
			"seed":        seed,
			"width":       w,
			"height":      h,
			"latency_ms":  time.Since(start).Milliseconds(),
		},
	}, nil
}

func blend(base uint8, f float64) uint8 {
	v := float64(base) + (255-float64(base))*f
	if v < 0 {
		v = 0
	}
	if v > 255 {
		v = 255
	}
	return uint8(v)
}

// derivedSeed produces a stable int64 seed from a prompt so tests are
// reproducible without forcing the caller to pick one.
func derivedSeed(prompt string) int64 {
	sum := sha256.Sum256([]byte(prompt))
	return int64(binary.BigEndian.Uint64(sum[:8]) >> 1)
}

func init() {
	Register("mock", NewMock)
}