easyai-ai-gateway/apps/api/internal/httpapi/gemini_compat_test.go

package httpapi

import (
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
)

func TestGeminiGenerateContentModelFromPath(t *testing.T) {
	tests := []struct {
		name        string
		prefix      string
		requestPath string
		wantModel   string
		wantOK      bool
	}{
		{
			name:        "v1beta model",
			prefix:      "/v1beta/models/",
			requestPath: "/v1beta/models/gemini-2.5-flash:generateContent",
			wantModel:   "gemini-2.5-flash",
			wantOK:      true,
		},
		{
			name:        "v1 model",
			prefix:      "/v1/models/",
			requestPath: "/v1/models/gemini-compatible:generateContent",
			wantModel:   "gemini-compatible",
			wantOK:      true,
		},
		{
			name:        "gateway api v1 model",
			prefix:      "/api/v1/models/",
			requestPath: "/api/v1/models/gemini-image:generateContent",
			wantModel:   "gemini-image",
			wantOK:      true,
		},
		{
			name:        "wrong suffix",
			prefix:      "/v1beta/models/",
			requestPath: "/v1beta/models/gemini-2.5-flash:countTokens",
			wantOK:      false,
		},
		{
			name:        "missing model",
			prefix:      "/v1beta/models/",
			requestPath: "/v1beta/models/:generateContent",
			wantOK:      false,
		},
		{
			name:        "extra path segment",
			prefix:      "/v1beta/models/",
			requestPath: "/v1beta/models/group/gemini-2.5-flash:generateContent",
			wantOK:      false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			gotModel, gotOK := geminiGenerateContentModelFromPath(tt.prefix, tt.requestPath)
			if gotOK != tt.wantOK || gotModel != tt.wantModel {
				t.Fatalf("geminiGenerateContentModelFromPath() = %q, %v; want %q, %v", gotModel, gotOK, tt.wantModel, tt.wantOK)
			}
		})
	}
}

func TestRegisterGeminiGenerateContentRoutes(t *testing.T) {
	server := &Server{
		auth:   auth.New("test-secret", "", ""),
		logger: slog.New(slog.NewTextHandler(io.Discard, nil)),
	}
	mux := http.NewServeMux()
	mux.HandleFunc("GET /api/v1/models", func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	})
	server.registerGeminiGenerateContentRoutes(mux)

	tests := []struct {
		method string
		path   string
		status int
	}{
		{method: http.MethodGet, path: "/api/v1/models", status: http.StatusNoContent},
		{method: http.MethodPost, path: "/api/v1/models/gemini-image:generateContent", status: http.StatusUnauthorized},
		{method: http.MethodPost, path: "/v1/models/gemini-image:generateContent", status: http.StatusUnauthorized},
		{method: http.MethodPost, path: "/v1beta/models/gemini-image:generateContent", status: http.StatusUnauthorized},
		{method: http.MethodPost, path: "/models/gemini-image:generateContent", status: http.StatusNotFound},
	}
	for _, tt := range tests {
		t.Run(tt.method+" "+tt.path, func(t *testing.T) {
			response := httptest.NewRecorder()
			mux.ServeHTTP(response, httptest.NewRequest(tt.method, tt.path, nil))
			if response.Code != tt.status {
				t.Fatalf("status = %d, want %d; body=%s", response.Code, tt.status, response.Body.String())
			}
		})
	}
}

func TestGeminiUploadPathKeepsCanonicalV1Prefix(t *testing.T) {
	canonical := httptest.NewRequest(http.MethodPost, "/api/v1/gemini/upload/v1beta/files", nil)
	if got := geminiUploadPath(canonical, "v1beta", "upload-1"); got != "/api/v1/gemini/upload/v1beta/files/upload-1" {
		t.Fatalf("canonical upload path = %q", got)
	}

	legacy := httptest.NewRequest(http.MethodPost, "/upload/v1beta/files", nil)
	if got := geminiUploadPath(legacy, "v1beta", "upload-1"); got != "/upload/v1beta/files/upload-1" {
		t.Fatalf("legacy upload path = %q", got)
	}
}

func TestGeminiImageTaskBodyMapsTextOnlyToImageGenerate(t *testing.T) {
	mapping, err := geminiImageTaskBody("gemini-image", map[string]any{
		"contents": []any{
			map[string]any{"parts": []any{
				map[string]any{"text": "draw"},
				map[string]any{"text": "a cat"},
			}},
		},
		"generationConfig": map[string]any{
			"candidateCount": float64(2),
			"imageConfig": map[string]any{
				"aspectRatio": "16:9",
				"imageSize":   "2K",
			},
		},
	})
	if err != nil {
		t.Fatalf("map Gemini request: %v", err)
	}
	if mapping.Kind != "images.generations" || mapping.Body["modelType"] != "image_generate" {
		t.Fatalf("unexpected mapping kind/body: %+v", mapping)
	}
	if mapping.Body["prompt"] != "draw\na cat" || mapping.Body["n"] != 2 || mapping.Body["aspect_ratio"] != "16:9" || mapping.Body["resolution"] != "2K" {
		t.Fatalf("unexpected mapped body: %+v", mapping.Body)
	}
}

func TestGeminiImageTaskBodyMapsInlineDataToImageEdit(t *testing.T) {
	mapping, err := geminiImageTaskBody("gemini-image", map[string]any{
		"contents": []any{
			map[string]any{"parts": []any{
				map[string]any{"text": "make it brighter"},
				map[string]any{"inlineData": map[string]any{
					"mimeType": "image/png",
					"data":     "aW1hZ2U=",
				}},
			}},
		},
	})
	if err != nil {
		t.Fatalf("map Gemini request: %v", err)
	}
	if mapping.Kind != "images.edits" || mapping.Body["modelType"] != "image_edit" {
		t.Fatalf("unexpected mapping kind/body: %+v", mapping)
	}
	if mapping.Body["prompt"] != "make it brighter" || mapping.Body["image"] != "data:image/png;base64,aW1hZ2U=" {
		t.Fatalf("unexpected mapped body: %+v", mapping.Body)
	}
}

func TestGeminiGenerateContentResponseMapsURLAndUsage(t *testing.T) {
	response := geminiGenerateContentResponse(map[string]any{
		"data": []any{map[string]any{
			"type": "image",
			"url":  "https://cdn.example/out.webp",
		}},
		"usage": map[string]any{
			"prompt_tokens":     3,
			"completion_tokens": 4,
			"total_tokens":      7,
		},
	}, "gemini-image")

	candidates := response["candidates"].([]any)
	content := candidates[0].(map[string]any)["content"].(map[string]any)
	parts := content["parts"].([]any)
	fileData := parts[0].(map[string]any)["fileData"].(map[string]any)
	if fileData["fileUri"] != "https://cdn.example/out.webp" || fileData["mimeType"] != "image/webp" {
		t.Fatalf("unexpected fileData part: %+v", fileData)
	}
	usage := response["usageMetadata"].(map[string]any)
	if usage["promptTokenCount"] != 3 || usage["candidatesTokenCount"] != 4 || usage["totalTokenCount"] != 7 {
		t.Fatalf("unexpected usage metadata: %+v", usage)
	}
}

func TestGeminiFileResponseShape(t *testing.T) {
	response := geminiFileResponse("files/test", "https://cdn.example/input.png", "image/png", 12)
	file := response["file"].(map[string]any)
	if file["name"] != "files/test" || file["uri"] != "https://cdn.example/input.png" || file["mimeType"] != "image/png" || file["sizeBytes"] != "12" || file["state"] != "ACTIVE" {
		t.Fatalf("unexpected file response: %+v", file)
	}
}