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

package httpapi

import "testing"

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:        "bare model path",
			prefix:      "/models/",
			requestPath: "/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 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)
	}
}