李建琦
6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
361 lines
14 KiB
Go
361 lines
14 KiB
Go
//go:build unit
|
|
|
|
package service
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"strings"
|
|
"testing"
|
|
|
|
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestBatchImageProviderRegistry_ReturnsGeminiAPI(t *testing.T) {
|
|
registry := NewDefaultBatchImageProviderRegistry()
|
|
provider, ok := registry.Get(BatchImageProviderGeminiAPI)
|
|
require.True(t, ok)
|
|
require.Equal(t, BatchImageProviderGeminiAPI, provider.Name())
|
|
|
|
must, err := registry.MustGet(BatchImageProviderGeminiAPI)
|
|
require.NoError(t, err)
|
|
require.Same(t, provider, must)
|
|
|
|
_, err = registry.MustGet("unknown_provider")
|
|
require.ErrorIs(t, err, ErrBatchImageInvalidProvider)
|
|
}
|
|
|
|
func TestGeminiProvider_SupportsOnlyGeminiAPIKeyWithSecret(t *testing.T) {
|
|
provider := NewGeminiAPIBatchImageProvider(&fakeGeminiBatchClient{})
|
|
|
|
require.True(t, provider.SupportsAccount(geminiAPIKeyAccount("sk-gemini")))
|
|
require.False(t, provider.SupportsAccount(&Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: map[string]any{}}))
|
|
require.False(t, provider.SupportsAccount(&Account{Platform: PlatformGemini, Type: AccountTypeOAuth, Credentials: map[string]any{"api_key": "sk"}}))
|
|
require.False(t, provider.SupportsAccount(&Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "sk"}}))
|
|
require.False(t, provider.SupportsAccount(nil))
|
|
}
|
|
|
|
func TestGeminiProvider_MissingAPIKeyRejected(t *testing.T) {
|
|
provider := NewGeminiAPIBatchImageProvider(&fakeGeminiBatchClient{})
|
|
_, err := provider.Submit(context.Background(), nil, &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey}, validGeminiBatchInput())
|
|
require.ErrorIs(t, err, ErrBatchImageProviderMissingAPIKey)
|
|
}
|
|
|
|
func TestBuildGeminiBatchJSONL_WritesValidLinesAndPreservesCustomID(t *testing.T) {
|
|
input := validGeminiBatchInput()
|
|
input.Items = append(input.Items, BatchImageInputItem{CustomID: "cover_002", Prompt: "Second prompt"})
|
|
|
|
jsonl, err := BuildGeminiBatchJSONL(input)
|
|
require.NoError(t, err)
|
|
|
|
lines := strings.Split(strings.TrimSpace(string(jsonl)), "\n")
|
|
require.Len(t, lines, 2)
|
|
requireJSONLLine(t, lines[0], "cover_001", "A clean product hero image")
|
|
requireJSONLLine(t, lines[1], "cover_002", "Second prompt")
|
|
}
|
|
|
|
func TestBuildGeminiBatchJSONL_RejectsDuplicateCustomIDs(t *testing.T) {
|
|
input := validGeminiBatchInput()
|
|
input.Items = append(input.Items, BatchImageInputItem{CustomID: "cover_001", Prompt: "Duplicate"})
|
|
|
|
_, err := BuildGeminiBatchJSONL(input)
|
|
require.ErrorIs(t, err, ErrBatchImageProviderInvalidInput)
|
|
}
|
|
|
|
func TestBuildGeminiBatchJSONL_RejectsEmptyPrompt(t *testing.T) {
|
|
input := validGeminiBatchInput()
|
|
input.Items[0].Prompt = " "
|
|
|
|
_, err := BuildGeminiBatchJSONL(input)
|
|
require.ErrorIs(t, err, ErrBatchImageProviderInvalidInput)
|
|
}
|
|
|
|
func TestBuildGeminiBatchJSONL_WritesReferenceImages(t *testing.T) {
|
|
input := validGeminiBatchInput()
|
|
input.Items[0].ReferenceImages = []BatchImageReference{
|
|
{MimeType: "image/webp", Data: []byte("webp-bytes")},
|
|
{MimeType: "image/jpeg", FileURI: "gs://bucket/refs/style.jpg"},
|
|
}
|
|
|
|
jsonl, err := BuildGeminiBatchJSONL(input)
|
|
require.NoError(t, err)
|
|
lines := strings.Split(strings.TrimSpace(string(jsonl)), "\n")
|
|
require.Len(t, lines, 1)
|
|
|
|
var got map[string]any
|
|
require.NoError(t, json.Unmarshal([]byte(lines[0]), &got))
|
|
request := got["request"].(map[string]any)
|
|
contents := request["contents"].([]any)
|
|
parts := contents[0].(map[string]any)["parts"].([]any)
|
|
require.Len(t, parts, 3)
|
|
require.Equal(t, "A clean product hero image", parts[0].(map[string]any)["text"])
|
|
inlineData := parts[1].(map[string]any)["inlineData"].(map[string]any)
|
|
require.Equal(t, "image/webp", inlineData["mimeType"])
|
|
require.Equal(t, "d2VicC1ieXRlcw==", inlineData["data"])
|
|
fileData := parts[2].(map[string]any)["fileData"].(map[string]any)
|
|
require.Equal(t, "image/jpeg", fileData["mimeType"])
|
|
require.Equal(t, "gs://bucket/refs/style.jpg", fileData["fileUri"])
|
|
}
|
|
|
|
func TestGeminiProvider_SubmitUploadsJSONLThenCreatesBatch(t *testing.T) {
|
|
client := &fakeGeminiBatchClient{
|
|
uploaded: &GeminiUploadedFile{Name: "files/input-jsonl"},
|
|
created: &GeminiBatchJob{Name: "batches/job-123", State: "JOB_STATE_PENDING"},
|
|
}
|
|
provider := NewGeminiAPIBatchImageProvider(client)
|
|
|
|
got, err := provider.Submit(context.Background(), &BatchImageJob{BatchID: "imgbatch_123", Model: "gemini-3.1-flash-image"}, geminiAPIKeyAccount("sk-secret"), validGeminiBatchInput())
|
|
require.NoError(t, err)
|
|
require.Equal(t, []string{"upload", "create"}, client.calls)
|
|
require.Equal(t, "files/input-jsonl", got.ProviderInputRef)
|
|
require.Equal(t, "batches/job-123", got.ProviderJobName)
|
|
require.Empty(t, got.ProviderOutputRef)
|
|
require.NotContains(t, got.ProviderInputRef, "A clean product hero image")
|
|
require.NotContains(t, string(client.uploadedJSONL), "sk-secret")
|
|
}
|
|
|
|
func TestGeminiProvider_GetMapsStates(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
job *GeminiBatchJob
|
|
wantState BatchProviderInternalState
|
|
wantDone bool
|
|
wantRef string
|
|
wantCode string
|
|
}{
|
|
{name: "running", job: &GeminiBatchJob{Name: "batches/1", State: "JOB_STATE_RUNNING"}, wantState: BatchProviderStateRunning},
|
|
{name: "succeeded_dest_fileName", job: &GeminiBatchJob{Name: "batches/1", State: "JOB_STATE_SUCCEEDED", Dest: &GeminiBatchDest{FileName: "files/out"}}, wantState: BatchProviderStateSucceeded, wantDone: true, wantRef: "files/out"},
|
|
{name: "failed", job: &GeminiBatchJob{Name: "batches/1", State: "JOB_STATE_FAILED", Error: &GeminiBatchError{Code: "BAD_PROMPT", Message: "bad prompt"}}, wantState: BatchProviderStateFailed, wantDone: true, wantCode: "BAD_PROMPT"},
|
|
{name: "cancelled", job: &GeminiBatchJob{Name: "batches/1", State: "JOB_STATE_CANCELLED"}, wantState: BatchProviderStateCancelled, wantDone: true, wantCode: "GEMINI_BATCH_CANCELLED"},
|
|
{name: "expired", job: &GeminiBatchJob{Name: "batches/1", State: "JOB_STATE_EXPIRED"}, wantState: BatchProviderStateExpired, wantDone: true, wantCode: "GEMINI_BATCH_EXPIRED"},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
provider := NewGeminiAPIBatchImageProvider(&fakeGeminiBatchClient{got: tt.job})
|
|
got, err := provider.Get(context.Background(), jobWithProviderName("batches/1"), geminiAPIKeyAccount("sk-secret"))
|
|
require.NoError(t, err)
|
|
require.Equal(t, tt.wantState, got.InternalState)
|
|
require.Equal(t, tt.wantDone, got.Done)
|
|
require.Equal(t, tt.wantRef, got.ProviderOutputRef)
|
|
require.Equal(t, tt.wantCode, got.ErrorCode)
|
|
require.NotContains(t, got.ErrorMessage, "sk-secret")
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestGeminiProvider_GetExtractsResponsesFileReference(t *testing.T) {
|
|
provider := NewGeminiAPIBatchImageProvider(&fakeGeminiBatchClient{
|
|
got: &GeminiBatchJob{
|
|
Name: "batches/1",
|
|
State: "JOB_STATE_SUCCEEDED",
|
|
Response: &GeminiBatchResponse{ResponsesFile: "files/responses-jsonl"},
|
|
},
|
|
})
|
|
|
|
got, err := provider.Get(context.Background(), jobWithProviderName("batches/1"), geminiAPIKeyAccount("sk-secret"))
|
|
require.NoError(t, err)
|
|
require.Equal(t, BatchProviderStateSucceeded, got.InternalState)
|
|
require.Equal(t, "files/responses-jsonl", got.ProviderOutputRef)
|
|
}
|
|
|
|
func TestGeminiProvider_GetRejectsInlineResultShape(t *testing.T) {
|
|
provider := NewGeminiAPIBatchImageProvider(&fakeGeminiBatchClient{
|
|
got: &GeminiBatchJob{
|
|
Name: "batches/1",
|
|
State: "JOB_STATE_SUCCEEDED",
|
|
Response: &GeminiBatchResponse{InlinedResponses: []any{map[string]any{"response": "large"}}},
|
|
},
|
|
})
|
|
|
|
_, err := provider.Get(context.Background(), jobWithProviderName("batches/1"), geminiAPIKeyAccount("sk-secret"))
|
|
require.ErrorIs(t, err, ErrBatchImageProviderInlineResultUnsupported)
|
|
}
|
|
|
|
func TestGeminiProvider_OpenResultStreamsResultFile(t *testing.T) {
|
|
client := &fakeGeminiBatchClient{downloadBody: "line1\n", downloadContentType: "application/jsonl"}
|
|
provider := NewGeminiAPIBatchImageProvider(client)
|
|
|
|
outputRef := "files/output-jsonl"
|
|
r, contentType, err := provider.OpenResult(context.Background(), &BatchImageJob{ProviderOutputRef: &outputRef}, geminiAPIKeyAccount("sk-secret"))
|
|
require.NoError(t, err)
|
|
defer r.Close()
|
|
|
|
body, err := io.ReadAll(r)
|
|
require.NoError(t, err)
|
|
require.Equal(t, "line1\n", string(body))
|
|
require.Equal(t, "application/jsonl", contentType)
|
|
require.Equal(t, "files/output-jsonl", client.downloadedFile)
|
|
}
|
|
|
|
func TestGeminiProvider_CancelCallsClient(t *testing.T) {
|
|
client := &fakeGeminiBatchClient{}
|
|
provider := NewGeminiAPIBatchImageProvider(client)
|
|
|
|
require.NoError(t, provider.Cancel(context.Background(), jobWithProviderName("batches/1"), geminiAPIKeyAccount("sk-secret")))
|
|
require.Equal(t, "batches/1", client.cancelledBatch)
|
|
}
|
|
|
|
func TestGeminiProvider_CleanupDeletesRefsOnlyWhenPresent(t *testing.T) {
|
|
inputRef := "files/input"
|
|
outputRef := "files/output"
|
|
client := &fakeGeminiBatchClient{}
|
|
provider := NewGeminiAPIBatchImageProvider(client)
|
|
|
|
err := provider.Cleanup(context.Background(), &BatchImageJob{ProviderInputRef: &inputRef, ProviderOutputRef: &outputRef}, geminiAPIKeyAccount("sk-secret"), CleanupTargetAll)
|
|
require.NoError(t, err)
|
|
require.Equal(t, []string{"files/input", "files/output"}, client.deletedFiles)
|
|
|
|
err = provider.Cleanup(context.Background(), &BatchImageJob{}, geminiAPIKeyAccount("sk-secret"), CleanupTargetAll)
|
|
require.NoError(t, err)
|
|
require.Equal(t, []string{"files/input", "files/output"}, client.deletedFiles)
|
|
}
|
|
|
|
func TestGeminiProvider_ErrorsDoNotExposeAPIKey(t *testing.T) {
|
|
apiKey := "sk-top-secret"
|
|
client := &fakeGeminiBatchClient{uploadErr: &GeminiAPIError{StatusCode: 401, Message: "upstream body should be hidden " + apiKey}}
|
|
provider := NewGeminiAPIBatchImageProvider(client)
|
|
|
|
_, err := provider.Submit(context.Background(), nil, geminiAPIKeyAccount(apiKey), validGeminiBatchInput())
|
|
require.Error(t, err)
|
|
require.Equal(t, "GEMINI_AUTH_FAILED", infraerrors.Reason(err))
|
|
require.NotContains(t, err.Error(), apiKey)
|
|
}
|
|
|
|
func TestGeminiProvider_MetadataDoesNotStoreImageBytesOrBase64(t *testing.T) {
|
|
client := &fakeGeminiBatchClient{
|
|
uploaded: &GeminiUploadedFile{Name: "files/input-jsonl"},
|
|
created: &GeminiBatchJob{Name: "batches/job-123", State: "JOB_STATE_PENDING"},
|
|
}
|
|
provider := NewGeminiAPIBatchImageProvider(client)
|
|
|
|
got, err := provider.Submit(context.Background(), nil, geminiAPIKeyAccount("sk-secret"), validGeminiBatchInput())
|
|
require.NoError(t, err)
|
|
require.NotContains(t, got.ProviderJobName, "base64")
|
|
require.NotContains(t, got.ProviderInputRef, "base64")
|
|
require.NotContains(t, got.ProviderOutputRef, "base64")
|
|
require.NotContains(t, got.ProviderJobName+got.ProviderInputRef+got.ProviderOutputRef, "iVBOR")
|
|
require.NotContains(t, got.ProviderJobName+got.ProviderInputRef+got.ProviderOutputRef, "A clean product hero image")
|
|
}
|
|
|
|
func requireJSONLLine(t *testing.T, line, wantKey, wantPrompt string) {
|
|
t.Helper()
|
|
var got map[string]any
|
|
require.NoError(t, json.Unmarshal([]byte(line), &got))
|
|
require.Equal(t, wantKey, got["key"])
|
|
request := got["request"].(map[string]any)
|
|
config := request["generationConfig"].(map[string]any)
|
|
require.Equal(t, []any{"TEXT", "IMAGE"}, config["responseModalities"])
|
|
contents := request["contents"].([]any)
|
|
parts := contents[0].(map[string]any)["parts"].([]any)
|
|
require.Equal(t, wantPrompt, parts[0].(map[string]any)["text"])
|
|
}
|
|
|
|
func validGeminiBatchInput() BatchImageInput {
|
|
return BatchImageInput{
|
|
BatchID: "imgbatch_123",
|
|
Model: "gemini-3.1-flash-image",
|
|
DisplayName: "test batch",
|
|
Items: []BatchImageInputItem{{
|
|
CustomID: "cover_001",
|
|
Prompt: "A clean product hero image",
|
|
}},
|
|
}
|
|
}
|
|
|
|
func geminiAPIKeyAccount(apiKey string) *Account {
|
|
return &Account{
|
|
Platform: PlatformGemini,
|
|
Type: AccountTypeAPIKey,
|
|
Credentials: map[string]any{"api_key": apiKey},
|
|
}
|
|
}
|
|
|
|
func jobWithProviderName(name string) *BatchImageJob {
|
|
return &BatchImageJob{ProviderJobName: &name}
|
|
}
|
|
|
|
type fakeGeminiBatchClient struct {
|
|
calls []string
|
|
uploaded *GeminiUploadedFile
|
|
created *GeminiBatchJob
|
|
got *GeminiBatchJob
|
|
uploadErr error
|
|
createErr error
|
|
getErr error
|
|
cancelErr error
|
|
downloadErr error
|
|
deleteErr error
|
|
uploadedJSONL []byte
|
|
createdFile string
|
|
cancelledBatch string
|
|
downloadedFile string
|
|
downloadBody string
|
|
downloadContentType string
|
|
deletedFiles []string
|
|
}
|
|
|
|
func (f *fakeGeminiBatchClient) UploadJSONL(_ context.Context, apiKey string, _ string, r io.Reader) (*GeminiUploadedFile, error) {
|
|
if strings.TrimSpace(apiKey) == "" {
|
|
return nil, errors.New("missing api key")
|
|
}
|
|
f.calls = append(f.calls, "upload")
|
|
f.uploadedJSONL, _ = io.ReadAll(r)
|
|
if f.uploadErr != nil {
|
|
return nil, f.uploadErr
|
|
}
|
|
if f.uploaded != nil {
|
|
return f.uploaded, nil
|
|
}
|
|
return &GeminiUploadedFile{Name: "files/input-jsonl"}, nil
|
|
}
|
|
|
|
func (f *fakeGeminiBatchClient) CreateBatch(_ context.Context, _ string, _ string, fileName string, _ string) (*GeminiBatchJob, error) {
|
|
f.calls = append(f.calls, "create")
|
|
f.createdFile = fileName
|
|
if f.createErr != nil {
|
|
return nil, f.createErr
|
|
}
|
|
if f.created != nil {
|
|
return f.created, nil
|
|
}
|
|
return &GeminiBatchJob{Name: "batches/job-123", State: "JOB_STATE_PENDING"}, nil
|
|
}
|
|
|
|
func (f *fakeGeminiBatchClient) GetBatch(_ context.Context, _ string, _ string) (*GeminiBatchJob, error) {
|
|
f.calls = append(f.calls, "get")
|
|
if f.getErr != nil {
|
|
return nil, f.getErr
|
|
}
|
|
return f.got, nil
|
|
}
|
|
|
|
func (f *fakeGeminiBatchClient) CancelBatch(_ context.Context, _ string, batchName string) error {
|
|
f.calls = append(f.calls, "cancel")
|
|
f.cancelledBatch = batchName
|
|
return f.cancelErr
|
|
}
|
|
|
|
func (f *fakeGeminiBatchClient) DownloadFile(_ context.Context, _ string, fileName string) (io.ReadCloser, string, error) {
|
|
f.calls = append(f.calls, "download")
|
|
f.downloadedFile = fileName
|
|
if f.downloadErr != nil {
|
|
return nil, "", f.downloadErr
|
|
}
|
|
contentType := f.downloadContentType
|
|
if contentType == "" {
|
|
contentType = "application/octet-stream"
|
|
}
|
|
return io.NopCloser(bytes.NewBufferString(f.downloadBody)), contentType, nil
|
|
}
|
|
|
|
func (f *fakeGeminiBatchClient) DeleteFile(_ context.Context, _ string, fileName string) error {
|
|
f.calls = append(f.calls, "delete")
|
|
f.deletedFiles = append(f.deletedFiles, fileName)
|
|
return f.deleteErr
|
|
}
|