yao/agent/llm/image_test.go

package llm

import (
	"encoding/base64"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestExtractImageFromResponse_B64(t *testing.T) {
	data := map[string]interface{}{
		"data": []interface{}{
			map[string]interface{}{
				"b64_json": "iVBORw0KGgoAAAANS...",
			},
		},
	}
	resp, err := extractImageFromResponse(data)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if resp.Image != "iVBORw0KGgoAAAANS..." {
		t.Errorf("got Image=%q, want %q", resp.Image, "iVBORw0KGgoAAAANS...")
	}
	if resp.Format != "png" {
		t.Errorf("got Format=%q, want %q", resp.Format, "png")
	}
}

func TestExtractImageFromResponse_URL(t *testing.T) {
	fakeImage := []byte{0xFF, 0xD8, 0xFF, 0xE0, 0x00, 0x10} // fake JPEG header bytes
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "image/jpeg")
		w.Write(fakeImage)
	}))
	defer srv.Close()

	data := map[string]interface{}{
		"data": []interface{}{
			map[string]interface{}{
				"b64_json": nil,
				"url":      srv.URL + "/image_0.jpeg",
			},
		},
	}
	resp, err := extractImageFromResponse(data)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	expected := base64.StdEncoding.EncodeToString(fakeImage)
	if resp.Image != expected {
		t.Errorf("got Image=%q, want %q", resp.Image, expected)
	}
	if resp.Format != "jpeg" {
		t.Errorf("got Format=%q, want %q", resp.Format, "jpeg")
	}
}

func TestExtractImageFromResponse_URLPng(t *testing.T) {
	fakeImage := []byte{0x89, 0x50, 0x4E, 0x47} // PNG magic bytes
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "image/png")
		w.Write(fakeImage)
	}))
	defer srv.Close()

	data := map[string]interface{}{
		"data": []interface{}{
			map[string]interface{}{
				"url": srv.URL + "/output.png",
			},
		},
	}
	resp, err := extractImageFromResponse(data)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if resp.Format != "png" {
		t.Errorf("got Format=%q, want %q", resp.Format, "png")
	}
	if resp.Image == "" {
		t.Error("expected non-empty base64 Image")
	}
}

func TestExtractImageFromResponse_URLDownloadFail(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusNotFound)
	}))
	defer srv.Close()

	data := map[string]interface{}{
		"data": []interface{}{
			map[string]interface{}{
				"url": srv.URL + "/missing.png",
			},
		},
	}
	_, err := extractImageFromResponse(data)
	if err == nil {
		t.Error("expected error for failed download")
	}
}

func TestExtractImageFromResponse_Empty(t *testing.T) {
	data := map[string]interface{}{
		"data": []interface{}{},
	}
	_, err := extractImageFromResponse(data)
	if err == nil {
		t.Error("expected error for empty data array")
	}
}

func TestExtractImageFromResponse_NoData(t *testing.T) {
	data := map[string]interface{}{}
	_, err := extractImageFromResponse(data)
	if err == nil {
		t.Error("expected error for missing data field")
	}
}

func TestExtractImageFromResponse_NullBoth(t *testing.T) {
	data := map[string]interface{}{
		"data": []interface{}{
			map[string]interface{}{
				"b64_json": nil,
				"url":      nil,
			},
		},
	}
	_, err := extractImageFromResponse(data)
	if err == nil {
		t.Error("expected error when both b64_json and url are null")
	}
}

func TestDownloadImageAsBase64(t *testing.T) {
	payload := []byte("fake-png-data")
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "image/png")
		w.Write(payload)
	}))
	defer srv.Close()

	b64, format, err := downloadImageAsBase64(srv.URL + "/test.png")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if format != "png" {
		t.Errorf("got format=%q, want %q", format, "png")
	}
	decoded, _ := base64.StdEncoding.DecodeString(b64)
	if string(decoded) != string(payload) {
		t.Errorf("decoded content mismatch")
	}
}

func TestDownloadImageAsBase64_FormatFromURL(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/octet-stream")
		w.Write([]byte("data"))
	}))
	defer srv.Close()

	_, format, err := downloadImageAsBase64(srv.URL + "/image.webp")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if format != "webp" {
		t.Errorf("got format=%q, want %q (from URL fallback)", format, "webp")
	}
}

func TestExtractAPIError_WithMessage(t *testing.T) {
	data := map[string]interface{}{
		"error": map[string]interface{}{
			"message": "insufficient quota",
		},
	}
	msg := extractAPIError(data)
	if msg != "insufficient quota" {
		t.Errorf("got %q, want %q", msg, "insufficient quota")
	}
}

func TestExtractAPIError_NoMessage(t *testing.T) {
	data := map[string]interface{}{
		"something": "else",
	}
	msg := extractAPIError(data)
	raw, _ := json.Marshal(data)
	if msg != string(raw) {
		t.Errorf("got %q, want raw JSON fallback", msg)
	}
}