1 files changed, 333 insertions, 0 deletions
diff --git a/internal/image/gemini_test.go b/internal/image/gemini_test.go
new file mode 100644
index 0000000..8c34c6e
--- /dev/null
+++ b/internal/image/gemini_test.go
@@ -0,0 +1,333 @@
+package image
+
+import (
+	"bytes"
+	"context"
+	"encoding/base64"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"google.golang.org/genai"
+)
+
+func TestNewGeminiProvider(t *testing.T) {
+	t.Parallel()
+
+	client := NewGeminiProvider(&GeminiConfig{APIKey: "test-key"})
+	if client == nil {
+		t.Fatal("expected client")
+	}
+	if client.config == nil {
+		t.Fatal("expected normalized config")
+	}
+	if client.config.Model != DefaultGeminiImageModel {
+		t.Fatalf("expected default model %q, got %q", DefaultGeminiImageModel, client.config.Model)
+	}
+	if client.config.TextModel != DefaultGeminiTextModel {
+		t.Fatalf("expected default text model %q, got %q", DefaultGeminiTextModel, client.config.TextModel)
+	}
+	if client.Name() != Gemini {
+		t.Fatalf("Name() = %q, want %q", client.Name(), Gemini)
+	}
+}
+
+func TestGeminiProvider_NoAPIKey(t *testing.T) {
+	client := NewGeminiProvider(&GeminiConfig{})
+
+	_, err := client.Search(context.Background(), DefaultSearchOptions("ябълка"))
+	if err == nil {
+		t.Fatal("expected error for missing API key")
+	}
+
+	searchErr, ok := err.(*SearchError)
+	if !ok {
+		t.Fatalf("expected SearchError, got %T", err)
+	}
+	if searchErr.Code != "NO_API_KEY" {
+		t.Fatalf("expected NO_API_KEY error, got %s", searchErr.Code)
+	}
+}
+
+func TestGeminiProvider_Search_CustomPromptSkipsTextGeneration(t *testing.T) {
+	originalText := geminiGenerateText
+	originalImage := geminiGenerateImage
+	t.Cleanup(func() {
+		geminiGenerateText = originalText
+		geminiGenerateImage = originalImage
+	})
+
+	geminiGenerateText = func(context.Context, *GeminiProvider, string, string, string, float32, int32) (string, error) {
+		t.Fatal("unexpected text generation for custom prompt")
+		return "", nil
+	}
+
+	var gotPrompt string
+	geminiGenerateImage = func(_ context.Context, _ *GeminiProvider, prompt, _ string) ([]byte, string, error) {
+		gotPrompt = prompt
+		return mustJPEGBytes(t), "image/jpeg", nil
+	}
+
+	client := NewGeminiProvider(&GeminiConfig{APIKey: "test-key"})
+	callbackCalled := false
+	client.SetPromptCallback(func(prompt string) {
+		callbackCalled = true
+		if prompt != "custom flashcard prompt" {
+			t.Fatalf("callback prompt = %q, want %q", prompt, "custom flashcard prompt")
+		}
+	})
+
+	results, err := client.Search(context.Background(), &SearchOptions{
+		Query:        "ябълка",
+		Translation:  "banana",
+		CustomPrompt: " custom flashcard prompt ",
+	})
+	if err != nil {
+		t.Fatalf("Search() unexpected error: %v", err)
+	}
+	if gotPrompt != "custom flashcard prompt" {
+		t.Fatalf("image prompt = %q, want %q", gotPrompt, "custom flashcard prompt")
+	}
+	if !callbackCalled {
+		t.Fatal("expected prompt callback to be called")
+	}
+	if client.LastPrompt() != "custom flashcard prompt" {
+		t.Fatalf("LastPrompt() = %q, want %q", client.LastPrompt(), "custom flashcard prompt")
+	}
+	if len(results) != 1 {
+		t.Fatalf("expected 1 result, got %d", len(results))
+	}
+	if results[0].Source != Gemini {
+		t.Fatalf("Source = %q, want %q", results[0].Source, Gemini)
+	}
+	if !strings.HasPrefix(results[0].URL, geminiDataPrefix) {
+		t.Fatalf("expected PNG data URI, got %q", results[0].URL)
+	}
+}
+
+func TestGeminiProvider_Search_GeneratedPromptFlow(t *testing.T) {
+	originalText := geminiGenerateText
+	originalImage := geminiGenerateImage
+	originalStyles := append([]string(nil), ArtisticStyles...)
+	t.Cleanup(func() {
+		geminiGenerateText = originalText
+		geminiGenerateImage = originalImage
+		ArtisticStyles = originalStyles
+	})
+
+	ArtisticStyles = []string{"Photorealism"}
+
+	var sceneCalls int
+	var gotPrompt string
+	var callbackPrompt string
+	geminiGenerateText = func(_ context.Context, _ *GeminiProvider, _, systemPrompt, userPrompt string, temperature float32, maxOutputTokens int32) (string, error) {
+		if strings.Contains(systemPrompt, "educational flashcards for language learning") {
+			sceneCalls++
+			if temperature != 0.7 || maxOutputTokens != 100 {
+				t.Fatalf("scene params = %v/%d, want 0.7/100", temperature, maxOutputTokens)
+			}
+			if !strings.Contains(userPrompt, "apple") {
+				t.Fatalf("scene prompt = %q, want English translation", userPrompt)
+			}
+			return "A bright apple sits centered on a wooden table.", nil
+		}
+		t.Fatalf("unexpected system prompt: %q", systemPrompt)
+		return "", nil
+	}
+
+	geminiGenerateImage = func(_ context.Context, _ *GeminiProvider, prompt, _ string) ([]byte, string, error) {
+		gotPrompt = prompt
+		return mustJPEGBytes(t), "image/jpeg", nil
+	}
+
+	client := NewGeminiProvider(&GeminiConfig{APIKey: "test-key"})
+	callbackCalled := false
+	client.SetPromptCallback(func(prompt string) {
+		callbackCalled = true
+		callbackPrompt = prompt
+	})
+
+	results, err := client.Search(context.Background(), &SearchOptions{
+		Query:       "ябълка",
+		Translation: "apple",
+	})
+	if err != nil {
+		t.Fatalf("Search() unexpected error: %v", err)
+	}
+	if sceneCalls != 1 {
+		t.Fatalf("sceneCalls = %d, want 1", sceneCalls)
+	}
+	if !callbackCalled {
+		t.Fatal("expected prompt callback to be called")
+	}
+	if callbackPrompt != gotPrompt {
+		t.Fatalf("prompt callback = %q, want %q", callbackPrompt, gotPrompt)
+	}
+	if client.LastPrompt() != gotPrompt {
+		t.Fatalf("LastPrompt() = %q, want %q", client.LastPrompt(), gotPrompt)
+	}
+	if len(results) != 1 {
+		t.Fatalf("expected 1 result, got %d", len(results))
+	}
+
+	result := results[0]
+	if result.Source != Gemini {
+		t.Fatalf("Source = %q, want %q", result.Source, Gemini)
+	}
+	if result.Width != 1 || result.Height != 1 {
+		t.Fatalf("Size = %dx%d, want %dx%d", result.Width, result.Height, 1, 1)
+	}
+	if !strings.Contains(result.Description, "apple") {
+		t.Fatalf("Description = %q, want translated word", result.Description)
+	}
+	if !strings.Contains(gotPrompt, "Generate a Photorealism educational flashcard image illustrating \"apple\".") {
+		t.Fatalf("Prompt = %q, want translated subject in generated prompt", gotPrompt)
+	}
+	if !strings.Contains(gotPrompt, "Scene: A bright apple sits centered on a wooden table.") {
+		t.Fatalf("Prompt = %q, want generated scene in prompt", gotPrompt)
+	}
+
+	reader, err := client.Download(context.Background(), result.URL)
+	if err != nil {
+		t.Fatalf("Download() unexpected error: %v", err)
+	}
+	t.Cleanup(func() {
+		_ = reader.Close()
+	})
+
+	data, err := io.ReadAll(reader)
+	if err != nil {
+		t.Fatalf("ReadAll() unexpected error: %v", err)
+	}
+	if !bytes.HasPrefix(data, []byte("\x89PNG\r\n\x1a\n")) {
+		t.Fatalf("Search output was not normalized to PNG")
+	}
+}
+
+func TestGeminiProvider_Search_InvalidOptions(t *testing.T) {
+	t.Parallel()
+
+	client := NewGeminiProvider(&GeminiConfig{APIKey: "test-key"})
+	_, err := client.Search(context.Background(), nil)
+	if err == nil {
+		t.Fatal("expected error for nil options")
+	}
+
+	searchErr, ok := err.(*SearchError)
+	if !ok {
+		t.Fatalf("expected SearchError, got %T", err)
+	}
+	if searchErr.Code != "INVALID_OPTIONS" {
+		t.Fatalf("expected INVALID_OPTIONS, got %s", searchErr.Code)
+	}
+}
+
+func TestGeminiProvider_DownloadDataURI(t *testing.T) {
+	client := &GeminiProvider{}
+	payload := mustPNGBytes(t)
+	url := geminiDataPrefix + base64.StdEncoding.EncodeToString(payload)
+
+	reader, err := client.Download(context.Background(), url)
+	if err != nil {
+		t.Fatalf("Download() unexpected error: %v", err)
+	}
+	t.Cleanup(func() {
+		_ = reader.Close()
+	})
+
+	data, err := io.ReadAll(reader)
+	if err != nil {
+		t.Fatalf("ReadAll() unexpected error: %v", err)
+	}
+	if !bytes.Equal(data, payload) {
+		t.Fatalf("Download() = %v, want %v", data, payload)
+	}
+}
+
+func TestGeminiProvider_DownloadHTTPFallback(t *testing.T) {
+	client := &GeminiProvider{}
+	payload := []byte("fallback image bytes")
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodGet {
+			t.Fatalf("request method = %s, want GET", r.Method)
+		}
+		_, _ = w.Write(payload)
+	}))
+	t.Cleanup(server.Close)
+
+	reader, err := client.Download(context.Background(), server.URL)
+	if err != nil {
+		t.Fatalf("Download() unexpected error: %v", err)
+	}
+	t.Cleanup(func() {
+		_ = reader.Close()
+	})
+
+	data, err := io.ReadAll(reader)
+	if err != nil {
+		t.Fatalf("ReadAll() unexpected error: %v", err)
+	}
+	if !bytes.Equal(data, payload) {
+		t.Fatalf("Download() = %v, want %v", data, payload)
+	}
+}
+
+func TestExtractGeneratedImage(t *testing.T) {
+	t.Parallel()
+
+	want := mustPNGBytes(t)
+	resp := &genai.GenerateContentResponse{
+		Candidates: []*genai.Candidate{
+			{
+				Content: &genai.Content{
+					Parts: []*genai.Part{
+						{InlineData: &genai.Blob{Data: want, MIMEType: "image/png"}},
+					},
+				},
+			},
+		},
+	}
+
+	got, mimeType, err := extractGeneratedImage(resp)
+	if err != nil {
+		t.Fatalf("extractGeneratedImage() error = %v", err)
+	}
+	if mimeType != "image/png" {
+		t.Fatalf("mimeType = %q, want %q", mimeType, "image/png")
+	}
+	if !bytes.Equal(got, want) {
+		t.Fatalf("extractGeneratedImage() = %v, want %v", got, want)
+	}
+}
+
+func TestEncodeAndNormalizePNG(t *testing.T) {
+	t.Parallel()
+
+	jpegBytes := mustJPEGBytes(t)
+	pngBytes, err := normalizePNG(jpegBytes, "image/jpeg")
+	if err != nil {
+		t.Fatalf("normalizePNG() unexpected error: %v", err)
+	}
+	if !bytes.HasPrefix(pngBytes, []byte("\x89PNG\r\n\x1a\n")) {
+		t.Fatalf("normalizePNG() did not return PNG data")
+	}
+
+	dataURL, err := encodeDataURL(jpegBytes, "image/jpeg")
+	if err != nil {
+		t.Fatalf("encodeDataURL() unexpected error: %v", err)
+	}
+	if !strings.HasPrefix(dataURL, geminiDataPrefix) {
+		t.Fatalf("encodeDataURL() = %q, want PNG data URI", dataURL)
+	}
+}
+
+func TestDecodeDataURIErrors(t *testing.T) {
+	t.Parallel()
+
+	if _, err := decodeDataURL("not-a-data-uri"); err == nil {
+		t.Fatal("expected error for invalid data URI")
+	}
+}