Cyrene_v2/internal/runtime/multimodal/vision/mime_test.go

package vision

import (
	"bytes"
	"encoding/base64"
	"errors"
	"image"
	"image/color"
	"image/png"
	"testing"
)

// pngBytes 返回一张真实的最小 PNG（1x1），使「字节嗅探」在测试里
// 有真实依据，而不是靠手写常量。
func pngBytes(t *testing.T) []byte {
	t.Helper()
	img := image.NewRGBA(image.Rect(0, 0, 1, 1))
	img.Set(0, 0, color.RGBA{R: 0x11, G: 0x22, B: 0x33, A: 0xff})
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		t.Fatalf("encode png: %v", err)
	}
	return buf.Bytes()
}

// htmlBytes 模拟网关/PROXY 的错误页：非图像响应体的典型形态。
func htmlBytes() []byte {
	return []byte("<!DOCTYPE html><html><body>502 Bad Gateway</body></html>")
}

// webpBytes 构造一个满足 RIFF/WEBPVP8 嗅探签名的 WebP 头。
func webpBytes() []byte {
	return append([]byte("RIFF\x00\x00\x00\x00WEBPVP8 "), make([]byte, 16)...)
}

// bmpBytes 构造一个满足 "BM" 嗅探签名的 BMP 头。
func bmpBytes() []byte {
	return append([]byte("BM"), make([]byte, 32)...)
}

// TestImageMIMEFor 覆盖「内容类型不能瞎猜」的裁决表：声明的类型只作参考，
// 真实字节嗅探是否决权，判不出图像一律报错（绝不静默回退 image/png）。
func TestImageMIMEFor(t *testing.T) {
	t.Parallel()
	tests := []struct {
		name     string
		declared string
		data     []byte
		wantMIME string
		wantErr  bool
	}{
		{name: "png with correct type", declared: "image/png", data: pngBytes(t), wantMIME: "image/png"},
		{name: "png with charset param", declared: "image/png; charset=binary", data: pngBytes(t), wantMIME: "image/png"},
		{name: "png with missing type sniffs", declared: "", data: pngBytes(t), wantMIME: "image/png"},
		{name: "bytes win over a wrong declaration", declared: "image/jpg", data: pngBytes(t), wantMIME: "image/png"},
		{name: "webp sniffs", declared: "", data: webpBytes(), wantMIME: "image/webp"},
		{name: "bmp sniffs", declared: "", data: bmpBytes(), wantMIME: "image/bmp"},
		{name: "html declared as html", declared: "text/html; charset=utf-8", data: htmlBytes(), wantErr: true},
		{name: "html with missing type", declared: "", data: htmlBytes(), wantErr: true},
		{name: "html falsely declared png", declared: "image/png", data: htmlBytes(), wantErr: true},
		{name: "svg rejected", declared: "image/svg+xml", data: []byte(`<svg xmlns="http://www.w3.org/2000/svg"/>`), wantErr: true},
		{name: "empty payload", declared: "", data: nil, wantErr: true},
		{name: "unknown binary without type", declared: "", data: []byte{0x00, 0x01, 0x02, 0x03}, wantErr: true},
		{name: "declared image with unsniffable bytes trusts the declaration",
			declared: "image/avif", data: []byte{0x00, 0x00, 0x00, 0x20}, wantMIME: "image/avif"},
		{name: "declared jpg alias normalises in the fallback",
			declared: "image/jpg", data: []byte{0x00, 0x01, 0x02, 0x03}, wantMIME: "image/jpeg"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			got, err := imageMIMEFor(tt.declared, tt.data)
			if tt.wantErr {
				if err == nil {
					t.Fatalf("imageMIMEFor(%q) = %q, want error", tt.declared, got)
				}
				if !errors.Is(err, ErrNotImage) {
					t.Fatalf("error = %v, want ErrNotImage", err)
				}
				return
			}
			if err != nil {
				t.Fatalf("imageMIMEFor(%q) error: %v", tt.declared, err)
			}
			if got != tt.wantMIME {
				t.Fatalf("mime = %q, want %q", got, tt.wantMIME)
			}
		})
	}
}

// TestMimeByExtLeavesUnknownEmpty 钉住未知扩展名不再被猜成 png —— 交给
// 字节嗅探裁决。
func TestMimeByExtLeavesUnknownEmpty(t *testing.T) {
	t.Parallel()
	if got := mimeByExt("/tmp/notes.txt"); got != "" {
		t.Fatalf("mimeByExt(.txt) = %q, want empty (let the bytes decide)", got)
	}
	if got := mimeByExt("/tmp/cat.JPEG"); got != "image/jpeg" {
		t.Fatalf("mimeByExt(.JPEG) = %q, want image/jpeg", got)
	}
	if got := mimeByExt("/tmp/shot.bmp"); got != "image/bmp" {
		t.Fatalf("mimeByExt(.bmp) = %q, want image/bmp", got)
	}
}

// TestReadFileRejectsRenamedHTML 钉住本地路径也不喂错类型：把 HTML 改名成
// .png 不能再被当成 PNG 读进模型。
func TestReadFileRejectsRenamedHTML(t *testing.T) {
	t.Parallel()
	path := t.TempDir() + "/liar.png"
	if err := writeFileBytes(path, htmlBytes()); err != nil {
		t.Fatalf("write: %v", err)
	}
	mime, data, err := readFile(path)
	if err == nil {
		t.Fatalf("readFile = (%q, %d bytes), want error", mime, len(data))
	}
	if !errors.Is(err, ErrNotImage) {
		t.Fatalf("error = %v, want ErrNotImage", err)
	}
}

// TestReadFileKeepsKnownImage 钉住正常本地图片不退化（含新增的 bmp）。
func TestReadFileKeepsKnownImage(t *testing.T) {
	t.Parallel()
	path := t.TempDir() + "/cat.png"
	want := pngBytes(t)
	if err := writeFileBytes(path, want); err != nil {
		t.Fatalf("write: %v", err)
	}
	mime, data, err := readFile(path)
	if err != nil {
		t.Fatalf("readFile: %v", err)
	}
	if mime != "image/png" {
		t.Fatalf("mime = %q, want image/png", mime)
	}
	if !bytes.Equal(data, want) {
		t.Fatal("bytes changed on read")
	}
}

// TestDataURIStable 是 base64 路径逐字节不变的底线。
func TestDataURIStable(t *testing.T) {
	t.Parallel()
	data := pngBytes(t)
	want := "data:image/png;base64," + base64.StdEncoding.EncodeToString(data)
	if got := DataURI("image/png", data); got != want {
		t.Fatalf("DataURI = %q, want %q", got, want)
	}
}