| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969 |
- //go:build unit
- package service
- import (
- "context"
- "encoding/json"
- "errors"
- "io"
- "strings"
- "testing"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
- "github.com/stretchr/testify/require"
- )
- func TestBatchImagePublicService_Submit(t *testing.T) {
- ctx := context.Background()
- t.Run("rejects when disabled", func(t *testing.T) {
- svc, _, _, _, _ := newTestBatchImagePublicService(false)
- _, err := svc.Submit(ctx, testBatchImageOwner(), validBatchImageSubmitRequest(), "")
- require.ErrorIs(t, err, ErrBatchImageDisabled)
- })
- t.Run("accepts valid request stores refs and enqueues once", func(t *testing.T) {
- svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true)
- req := validBatchImageSubmitRequest()
- req.SessionID = batchImageStringPtr("batch-session-123")
- got, err := svc.Submit(ctx, testBatchImageOwner(), req, "")
- require.NoError(t, err)
- require.Equal(t, "image.batch", got.Object)
- require.Equal(t, "queued", got.Status)
- require.Equal(t, BatchImageProviderGeminiAPI, got.Provider)
- require.Equal(t, 2, got.ItemCount)
- require.Equal(t, 0.25, got.EstimatedCost)
- require.Len(t, repo.jobs, 1)
- require.Len(t, gemini.submits, 1)
- require.Equal(t, []string{got.ID}, queue.enqueued)
- billing := svc.BillingRepo.(*fakeBatchImageBillingRepo)
- require.Len(t, billing.reserves, 1)
- require.Equal(t, BatchImageHoldRequestID(got.ID), billing.reserves[0].RequestID)
- require.InDelta(t, 0.3, billing.reserves[0].HoldAmount, 1e-12)
- require.Empty(t, billing.releases)
- authCache := svc.AuthCache.(*fakeBatchImageAuthCacheInvalidator)
- require.Equal(t, []int64{11}, authCache.userIDs)
- job := repo.jobs[got.ID]
- require.Equal(t, BatchImageJobStatusSubmitted, job.Status)
- require.Equal(t, "providers/gemini_api/job", batchImageDerefString(job.ProviderJobName))
- require.Equal(t, "files/gemini_api/input", batchImageDerefString(job.ProviderInputRef))
- require.Equal(t, "files/gemini_api/output", batchImageDerefString(job.ProviderOutputRef))
- require.NotNil(t, job.AccountID)
- require.Equal(t, int64(202), *job.AccountID)
- require.Equal(t, 1, job.PricingSnapshotVersion)
- require.InDelta(t, 0.25, job.BaseUnitPrice, 1e-12)
- require.InDelta(t, 1.0, job.GroupRateMultiplier, 1e-12)
- require.InDelta(t, 1.0, job.AccountRateMultiplier, 1e-12)
- require.InDelta(t, 0.5, job.BatchDiscountMultiplier, 1e-12)
- require.InDelta(t, 0.6, job.HoldMultiplier, 1e-12)
- require.InDelta(t, 0.125, job.BillableUnitPrice, 1e-12)
- require.InDelta(t, 0.15, job.HoldUnitPrice, 1e-12)
- require.Equal(t, "batch-session-123", batchImageDerefString(job.SessionID))
- })
- t.Run("combines user group image rate account rate discount and hold margin", func(t *testing.T) {
- svc, repo, _, _, _ := newTestBatchImagePublicService(true)
- groupID := int64(7)
- accountMultiplier := 1.25
- accountRepo := svc.AccountRepo.(*publicBatchImageAccountRepo)
- accountRepo.accounts[1].RateMultiplier = &accountMultiplier
- svc.GroupRepo = &publicBatchImageGroupRepo{groups: map[int64]*Group{
- groupID: {
- ID: groupID,
- Platform: PlatformGemini,
- RateMultiplier: 2.0,
- AllowImageGeneration: true,
- AllowBatchImageGeneration: true,
- ImageRateIndependent: false,
- BatchImageDiscountMultiplier: 0.8,
- BatchImageHoldMultiplier: 0.6,
- },
- }}
- userRate := 0.5
- svc.UserGroupRateRepo = &publicBatchImageUserGroupRateRepo{rates: map[int64]*float64{groupID: &userRate}}
- got, err := svc.Submit(ctx, BatchImageOwner{UserID: 11, APIKeyID: 22, GroupID: &groupID}, validBatchImageSubmitRequest(), "")
- require.NoError(t, err)
- require.InDelta(t, 0.25, got.EstimatedCost, 1e-12)
- job := repo.jobs[got.ID]
- require.InDelta(t, 0.25, job.BaseUnitPrice, 1e-12)
- require.InDelta(t, 0.5, job.GroupRateMultiplier, 1e-12)
- require.InDelta(t, 1.25, job.AccountRateMultiplier, 1e-12)
- require.InDelta(t, 0.8, job.BatchDiscountMultiplier, 1e-12)
- // 配置的 hold(0.6) < discount(0.8) 属于会导致结算死锁的脏数据,
- // 快照时被钳制为 discount,保证 holdAmount >= 实际成本上限。
- require.InDelta(t, 0.8, job.HoldMultiplier, 1e-12)
- require.InDelta(t, 0.125, job.BillableUnitPrice, 1e-12)
- require.InDelta(t, 0.125, job.HoldUnitPrice, 1e-12)
- require.InDelta(t, 0.25, *job.HoldAmount, 1e-12)
- })
- t.Run("uses configured group 1k image price for batch image base price", func(t *testing.T) {
- svc, repo, _, _, _ := newTestBatchImagePublicService(true)
- groupID := int64(7)
- imagePrice := 0.134
- svc.GroupRepo = &publicBatchImageGroupRepo{groups: map[int64]*Group{
- groupID: {
- ID: groupID,
- Platform: PlatformGemini,
- RateMultiplier: 1.0,
- AllowImageGeneration: true,
- AllowBatchImageGeneration: true,
- ImagePrice1K: &imagePrice,
- BatchImageDiscountMultiplier: 0.5,
- BatchImageHoldMultiplier: 0.6,
- },
- }}
- got, err := svc.Submit(ctx, BatchImageOwner{UserID: 11, APIKeyID: 22, GroupID: &groupID}, validBatchImageSubmitRequest(), "")
- require.NoError(t, err)
- require.InDelta(t, 0.134, got.EstimatedCost, 1e-12)
- job := repo.jobs[got.ID]
- require.InDelta(t, 0.134, job.BaseUnitPrice, 1e-12)
- require.InDelta(t, 0.067, job.BillableUnitPrice, 1e-12)
- require.InDelta(t, 0.0804, job.HoldUnitPrice, 1e-12)
- require.InDelta(t, 0.1608, *job.HoldAmount, 1e-12)
- })
- t.Run("pricing missing rejects before provider submit", func(t *testing.T) {
- svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true)
- svc.Pricing = &fakeBatchImagePricingResolver{err: ErrBatchImageSettlementPricingMissing}
- _, err := svc.Submit(ctx, testBatchImageOwner(), validBatchImageSubmitRequest(), "")
- require.ErrorIs(t, err, ErrBatchImageSettlementPricingMissing)
- require.Empty(t, repo.jobs)
- require.Empty(t, queue.enqueued)
- require.Empty(t, gemini.submits)
- })
- t.Run("group batch image disabled rejects before provider submit", func(t *testing.T) {
- svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true)
- groupID := int64(7)
- svc.GroupRepo = &publicBatchImageGroupRepo{groups: map[int64]*Group{
- groupID: {
- ID: groupID,
- Platform: PlatformGemini,
- RateMultiplier: 1,
- AllowBatchImageGeneration: false,
- BatchImageDiscountMultiplier: 0.5,
- BatchImageHoldMultiplier: 0.6,
- },
- }}
- _, err := svc.Submit(ctx, BatchImageOwner{UserID: 11, APIKeyID: 22, GroupID: &groupID}, validBatchImageSubmitRequest(), "")
- require.ErrorIs(t, err, ErrBatchImageGroupDisabled)
- require.Empty(t, repo.jobs)
- require.Empty(t, queue.enqueued)
- require.Empty(t, gemini.submits)
- })
- t.Run("group pricing load failure rejects before provider submit", func(t *testing.T) {
- svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true)
- groupID := int64(404)
- _, err := svc.Submit(ctx, BatchImageOwner{UserID: 11, APIKeyID: 22, GroupID: &groupID}, validBatchImageSubmitRequest(), "")
- require.ErrorIs(t, err, ErrBatchImageSettlementPricingMissing)
- require.Empty(t, repo.jobs)
- require.Empty(t, queue.enqueued)
- require.Empty(t, gemini.submits)
- })
- t.Run("generates custom ids deterministically", func(t *testing.T) {
- svc, _, _, gemini, _ := newTestBatchImagePublicService(true)
- req := validBatchImageSubmitRequest()
- req.Items[0].CustomID = ""
- req.Items[1].CustomID = ""
- _, err := svc.Submit(ctx, testBatchImageOwner(), req, "")
- require.NoError(t, err)
- require.Len(t, gemini.submits, 1)
- require.Equal(t, "item_000001", gemini.submits[0].Items[0].CustomID)
- require.Equal(t, "item_000002", gemini.submits[0].Items[1].CustomID)
- })
- t.Run("expands output count into separate billable items", func(t *testing.T) {
- svc, repo, _, gemini, _ := newTestBatchImagePublicService(true)
- req := validBatchImageSubmitRequest()
- req.Items = []BatchImageSubmitItem{
- {CustomID: "cover", Prompt: "hero", OutputCount: 3, ReferenceImages: []BatchImageReferenceInput{{MimeType: "image/png", Data: []byte("ref")}}},
- }
- got, err := svc.Submit(ctx, testBatchImageOwner(), req, "")
- require.NoError(t, err)
- require.Equal(t, 3, got.ItemCount)
- require.InDelta(t, 0.375, got.EstimatedCost, 1e-12)
- require.Len(t, gemini.submits, 1)
- require.Len(t, gemini.submits[0].Items, 3)
- require.Equal(t, []string{"cover_01", "cover_02", "cover_03"}, []string{
- gemini.submits[0].Items[0].CustomID,
- gemini.submits[0].Items[1].CustomID,
- gemini.submits[0].Items[2].CustomID,
- })
- require.Len(t, gemini.submits[0].Items[0].ReferenceImages, 1)
- require.Len(t, repo.items[got.ID], 3)
- })
- t.Run("validates request fields", func(t *testing.T) {
- tests := []struct {
- name string
- mutate func(*BatchImageSubmitRequest)
- want error
- }{
- {name: "missing_model", mutate: func(r *BatchImageSubmitRequest) { r.Model = "" }, want: ErrBatchImageInvalidModel},
- {name: "empty_items", mutate: func(r *BatchImageSubmitRequest) { r.Items = nil }, want: ErrBatchImageInvalidItems},
- {name: "duplicate_custom_ids", mutate: func(r *BatchImageSubmitRequest) { r.Items[1].CustomID = r.Items[0].CustomID }, want: ErrBatchImageDuplicateCustomIDInRequest},
- {name: "empty_prompt", mutate: func(r *BatchImageSubmitRequest) { r.Items[0].Prompt = " " }, want: ErrBatchImageInvalidItems},
- {name: "prompt_too_long", mutate: func(r *BatchImageSubmitRequest) { r.Items[0].Prompt = strings.Repeat("x", 9) }, want: ErrBatchImagePromptTooLong},
- {name: "unsupported_provider", mutate: func(r *BatchImageSubmitRequest) { r.Provider = "other" }, want: ErrBatchImageUnsupportedProvider},
- {name: "vertex_rejects_2k", mutate: func(r *BatchImageSubmitRequest) { r.Provider = BatchImageProviderVertex; r.ImageSize = "2K" }, want: ErrBatchImageInvalidItems},
- {name: "too_many_outputs_per_item", mutate: func(r *BatchImageSubmitRequest) {
- r.Items[0].OutputCount = 5
- }, want: ErrBatchImageInvalidItems},
- {name: "too_many_reference_images_for_flash", mutate: func(r *BatchImageSubmitRequest) {
- r.Model = "gemini-2.5-flash-image"
- r.Items[0].ReferenceImages = []BatchImageReferenceInput{
- {MimeType: "image/png", Data: []byte("1")},
- {MimeType: "image/png", Data: []byte("2")},
- {MimeType: "image/png", Data: []byte("3")},
- {MimeType: "image/png", Data: []byte("4")},
- }
- }, want: ErrBatchImageTooManyReferenceImages},
- {name: "bad_reference_mime", mutate: func(r *BatchImageSubmitRequest) {
- r.Items[0].ReferenceImages = []BatchImageReferenceInput{{MimeType: "application/octet-stream", Data: []byte("x")}}
- }, want: ErrBatchImageInvalidReferenceImage},
- {name: "reference_requires_data_or_file_uri", mutate: func(r *BatchImageSubmitRequest) {
- r.Items[0].ReferenceImages = []BatchImageReferenceInput{{MimeType: "image/png"}}
- }, want: ErrBatchImageInvalidReferenceImage},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- svc, _, _, _, _ := newTestBatchImagePublicService(true)
- req := validBatchImageSubmitRequest()
- tt.mutate(&req)
- _, err := svc.Submit(ctx, testBatchImageOwner(), req, "")
- require.ErrorIs(t, err, tt.want)
- })
- }
- })
- t.Run("rejects too many items", func(t *testing.T) {
- svc, _, _, _, _ := newTestBatchImagePublicService(true)
- req := validBatchImageSubmitRequest()
- req.Items = append(req.Items, BatchImageSubmitItem{CustomID: "too_many", Prompt: "x"})
- _, err := svc.Submit(ctx, testBatchImageOwner(), req, "")
- require.ErrorIs(t, err, ErrBatchImageInvalidItems)
- })
- t.Run("rejects too many output images", func(t *testing.T) {
- svc, _, _, _, _ := newTestBatchImagePublicService(true)
- svc.Config.BatchImage.MaxOutputImagesPerJob = 3
- req := validBatchImageSubmitRequest()
- req.Items[0].OutputCount = 2
- req.Items[1].OutputCount = 2
- _, err := svc.Submit(ctx, testBatchImageOwner(), req, "")
- require.ErrorIs(t, err, ErrBatchImageTooManyOutputImages)
- })
- t.Run("rejects too many reference images across request", func(t *testing.T) {
- svc, _, _, _, _ := newTestBatchImagePublicService(true)
- svc.Config.BatchImage.MaxReferenceImagesPerJob = 3
- req := validBatchImageSubmitRequest()
- req.Model = "gemini-2.5-flash-image"
- req.Items[0].ReferenceImages = []BatchImageReferenceInput{
- {MimeType: "image/png", Data: []byte("1")},
- {MimeType: "image/png", Data: []byte("2")},
- }
- req.Items[1].ReferenceImages = []BatchImageReferenceInput{
- {MimeType: "image/png", Data: []byte("3")},
- {MimeType: "image/png", Data: []byte("4")},
- }
- _, err := svc.Submit(ctx, testBatchImageOwner(), req, "")
- require.ErrorIs(t, err, ErrBatchImageTooManyReferenceImages)
- })
- t.Run("rejects too much inline reference image data across request", func(t *testing.T) {
- svc, _, _, _, _ := newTestBatchImagePublicService(true)
- svc.Config.BatchImage.MaxReferenceImagesPerJob = 10
- svc.Config.BatchImage.MaxReferenceInlineBytesPerJob = 4
- req := validBatchImageSubmitRequest()
- req.Model = "gemini-2.5-flash-image"
- req.Items[0].ReferenceImages = []BatchImageReferenceInput{{MimeType: "image/png", Data: []byte("123")}}
- req.Items[1].ReferenceImages = []BatchImageReferenceInput{{MimeType: "image/png", Data: []byte("456")}}
- _, err := svc.Submit(ctx, testBatchImageOwner(), req, "")
- require.ErrorIs(t, err, ErrBatchImageReferenceImagesTooLarge)
- })
- t.Run("selects requested provider", func(t *testing.T) {
- svc, _, _, gemini, vertex := newTestBatchImagePublicService(true)
- req := validBatchImageSubmitRequest()
- req.Provider = BatchImageProviderVertex
- got, err := svc.Submit(ctx, testBatchImageOwner(), req, "")
- require.NoError(t, err)
- require.Equal(t, BatchImageProviderVertex, got.Provider)
- require.Empty(t, gemini.submits)
- require.Len(t, vertex.submits, 1)
- })
- t.Run("insufficient balance rejects before provider submit", func(t *testing.T) {
- svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true)
- billing := &fakeBatchImageBillingRepo{err: ErrBatchImageInsufficientBalance}
- svc.BillingRepo = billing
- _, err := svc.Submit(ctx, testBatchImageOwner(), validBatchImageSubmitRequest(), "")
- require.ErrorIs(t, err, ErrBatchImageInsufficientBalance)
- require.Empty(t, queue.enqueued)
- require.Empty(t, gemini.submits)
- require.Len(t, billing.reserves, 1)
- require.Empty(t, billing.releases)
- require.Len(t, repo.jobs, 1)
- for _, job := range repo.jobs {
- require.Equal(t, BatchImageJobStatusFailed, job.Status)
- require.Equal(t, "INSUFFICIENT_BALANCE", batchImageDerefString(job.LastErrorCode))
- require.NotNil(t, job.UserDeletedAt)
- }
- })
- t.Run("provider failure marks failed and does not enqueue", func(t *testing.T) {
- svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true)
- gemini.submitErr = errors.New("projects/secret-provider-job failed")
- billing := svc.BillingRepo.(*fakeBatchImageBillingRepo)
- _, err := svc.Submit(ctx, testBatchImageOwner(), validBatchImageSubmitRequest(), "")
- require.ErrorIs(t, err, ErrBatchImageProviderSubmitFailed)
- require.Empty(t, queue.enqueued)
- require.Len(t, billing.reserves, 1)
- require.Len(t, billing.releases, 1)
- require.Equal(t, BatchImageReleaseRequestID(billing.reserves[0].BatchID), billing.releases[0].RequestID)
- require.Len(t, repo.jobs, 1)
- for _, job := range repo.jobs {
- require.Equal(t, BatchImageJobStatusFailed, job.Status)
- require.Equal(t, "PROVIDER_SUBMIT_FAILED", batchImageDerefString(job.LastErrorCode))
- require.Equal(t, "upstream provider operation failed", batchImageDerefString(job.LastErrorMessage))
- require.NotNil(t, job.UserDeletedAt)
- }
- })
- t.Run("provider failure with release failure enqueues billing retry", func(t *testing.T) {
- svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true)
- gemini.submitErr = errors.New("projects/secret-provider-job failed")
- billing := svc.BillingRepo.(*fakeBatchImageBillingRepo)
- billing.releaseErr = errors.New("billing database timeout")
- _, err := svc.Submit(ctx, testBatchImageOwner(), validBatchImageSubmitRequest(), "")
- require.ErrorIs(t, err, ErrBatchImageBillingHoldFailed)
- require.Len(t, billing.reserves, 1)
- require.Len(t, billing.releases, 1)
- require.Len(t, repo.jobs, 1)
- for _, job := range repo.jobs {
- require.Equal(t, BatchImageJobStatusFailed, job.Status)
- require.Equal(t, "BILLING_RELEASE_FAILED", batchImageDerefString(job.LastErrorCode))
- require.Equal(t, []string{job.BatchID}, queue.enqueued)
- }
- })
- t.Run("queue failure is recorded after provider submit", func(t *testing.T) {
- svc, repo, queue, _, _ := newTestBatchImagePublicService(true)
- queue.err = errors.New("redis unavailable")
- billing := svc.BillingRepo.(*fakeBatchImageBillingRepo)
- _, err := svc.Submit(ctx, testBatchImageOwner(), validBatchImageSubmitRequest(), "")
- require.ErrorIs(t, err, ErrBatchImageQueueFailed)
- require.Len(t, billing.reserves, 1)
- require.Empty(t, billing.releases)
- require.Len(t, repo.jobs, 1)
- for _, job := range repo.jobs {
- require.Equal(t, BatchImageJobStatusSubmitted, job.Status)
- require.Equal(t, "QUEUE_FAILED", batchImageDerefString(job.LastErrorCode))
- require.Contains(t, repo.events[job.BatchID], "queue_failed")
- }
- })
- t.Run("idempotency returns same batch without provider resubmit", func(t *testing.T) {
- svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true)
- req := validBatchImageSubmitRequest()
- req.SessionID = batchImageStringPtr("original-session")
- first, err := svc.Submit(ctx, testBatchImageOwner(), req, "client-key")
- require.NoError(t, err)
- req.SessionID = batchImageStringPtr("retry-session")
- second, err := svc.Submit(ctx, testBatchImageOwner(), req, "client-key")
- require.NoError(t, err)
- require.Equal(t, first.ID, second.ID)
- require.Equal(t, "original-session", batchImageDerefString(repo.jobs[first.ID].SessionID))
- require.Len(t, gemini.submits, 1)
- require.Equal(t, []string{first.ID}, queue.enqueued)
- })
- t.Run("idempotency conflict rejects changed request", func(t *testing.T) {
- svc, _, _, _, _ := newTestBatchImagePublicService(true)
- req := validBatchImageSubmitRequest()
- first, err := svc.Submit(ctx, testBatchImageOwner(), req, "client-key")
- require.NoError(t, err)
- req.Items[0].Prompt = "diff"
- second, err := svc.Submit(ctx, testBatchImageOwner(), req, "client-key")
- require.Nil(t, second)
- require.ErrorIs(t, err, ErrBatchImageIdempotencyConflict)
- require.NotEmpty(t, first.ID)
- })
- t.Run("public response does not expose internals", func(t *testing.T) {
- svc, _, _, _, _ := newTestBatchImagePublicService(true)
- got, err := svc.Submit(ctx, testBatchImageOwner(), validBatchImageSubmitRequest(), "")
- require.NoError(t, err)
- body, err := json.Marshal(got)
- require.NoError(t, err)
- requireBatchImagePublicJSONHasNoInternals(t, string(body))
- })
- }
- func TestBatchImagePublicService_List(t *testing.T) {
- ctx := context.Background()
- svc, repo, _, _, _ := newTestBatchImagePublicService(true)
- visibleKeyID := int64(22)
- otherKeyID := int64(23)
- repo.jobs["visible-1"] = &BatchImageJob{
- BatchID: "visible-1",
- UserID: 11,
- APIKeyID: &visibleKeyID,
- Status: BatchImageJobStatusCompleted,
- Provider: BatchImageProviderVertex,
- Model: "gemini-3.1-flash-lite-image",
- ItemCount: 1,
- CreatedAt: time.Now(),
- }
- repo.jobs["hidden-other-key"] = &BatchImageJob{
- BatchID: "hidden-other-key",
- UserID: 11,
- APIKeyID: &otherKeyID,
- Status: BatchImageJobStatusCompleted,
- Provider: BatchImageProviderVertex,
- Model: "gemini-3.1-flash-lite-image",
- ItemCount: 1,
- CreatedAt: time.Now(),
- }
- got, err := svc.List(ctx, BatchImageOwner{UserID: 11, APIKeyID: visibleKeyID}, BatchImageJobsQuery{Limit: 20})
- require.NoError(t, err)
- require.Equal(t, "list", got.Object)
- require.Len(t, got.Data, 1)
- require.Equal(t, "visible-1", got.Data[0].ID)
- require.False(t, got.HasMore)
- }
- func TestBatchImagePublicService_ListModels(t *testing.T) {
- ctx := context.Background()
- t.Run("requires explicit account model mapping", func(t *testing.T) {
- svc, _, _, _, _ := newTestBatchImagePublicService(true)
- got, err := svc.ListModels(ctx, testBatchImageOwner())
- require.NoError(t, err)
- require.Equal(t, "list", got.Object)
- require.Empty(t, got.Data)
- })
- t.Run("returns priced models from selected account group", func(t *testing.T) {
- svc, _, _, _, _ := newTestBatchImagePublicService(true)
- groupID := int64(7)
- svc.GroupRepo = &publicBatchImageGroupRepo{groups: map[int64]*Group{
- groupID: {
- ID: groupID,
- Platform: PlatformGemini,
- RateMultiplier: 1,
- AllowImageGeneration: true,
- AllowBatchImageGeneration: true,
- BatchImageDiscountMultiplier: 0.5,
- BatchImageHoldMultiplier: 0.6,
- },
- }}
- accountRepo := svc.AccountRepo.(*publicBatchImageAccountRepo)
- accountRepo.accounts = []Account{testBatchImageMappedAccount(303, AccountTypeAPIKey, map[string]any{
- "gemini-2.5-flash-image": "gemini-2.5-flash-image",
- })}
- got, err := svc.ListModels(ctx, BatchImageOwner{UserID: 11, APIKeyID: 22, GroupID: &groupID})
- require.NoError(t, err)
- require.Equal(t, []BatchImagePublicModel{{
- ID: "gemini-2.5-flash-image",
- Object: "image.batch.model",
- Provider: BatchImageProviderGeminiAPI,
- }, {
- ID: "gemini-2.5-flash-image",
- Object: "image.batch.model",
- Provider: BatchImageProviderVertex,
- }}, got.Data)
- })
- t.Run("expands wildcard mappings against batch image candidates", func(t *testing.T) {
- svc, _, _, _, _ := newTestBatchImagePublicService(true)
- accountRepo := svc.AccountRepo.(*publicBatchImageAccountRepo)
- accountRepo.accounts = []Account{testBatchImageMappedAccount(303, AccountTypeAPIKey, map[string]any{
- "gemini-3.1-*": "gemini-3.1-flash-lite-image",
- })}
- got, err := svc.ListModels(ctx, testBatchImageOwner())
- require.NoError(t, err)
- require.NotEmpty(t, got.Data)
- ids := make([]string, 0, len(got.Data))
- for _, model := range got.Data {
- ids = append(ids, model.ID)
- }
- require.Contains(t, ids, "gemini-3.1-flash-image")
- require.Contains(t, ids, "gemini-3.1-flash-lite-image")
- require.NotContains(t, ids, "gemini-2.5-flash-image")
- })
- t.Run("filters models without batch image pricing", func(t *testing.T) {
- svc, _, _, _, _ := newTestBatchImagePublicService(true)
- svc.Pricing = &fakeBatchImagePricingResolver{
- unitPrice: 0.25,
- missingModels: map[string]bool{"gemini-3.1-flash-lite-image": true},
- }
- accountRepo := svc.AccountRepo.(*publicBatchImageAccountRepo)
- accountRepo.accounts = []Account{testBatchImageMappedAccount(303, AccountTypeAPIKey, map[string]any{
- "gemini-2.5-flash-image": "gemini-2.5-flash-image",
- "gemini-3.1-flash-lite-image": "gemini-3.1-flash-lite-image",
- })}
- got, err := svc.ListModels(ctx, testBatchImageOwner())
- require.NoError(t, err)
- ids := make([]string, 0, len(got.Data))
- for _, model := range got.Data {
- ids = append(ids, model.ID)
- }
- require.Contains(t, ids, "gemini-2.5-flash-image")
- require.NotContains(t, ids, "gemini-3.1-flash-lite-image")
- })
- t.Run("rejects when group disables batch image", func(t *testing.T) {
- svc, _, _, _, _ := newTestBatchImagePublicService(true)
- groupID := int64(7)
- svc.GroupRepo = &publicBatchImageGroupRepo{groups: map[int64]*Group{
- groupID: {ID: groupID, AllowBatchImageGeneration: false},
- }}
- _, err := svc.ListModels(ctx, BatchImageOwner{UserID: 11, APIKeyID: 22, GroupID: &groupID})
- require.ErrorIs(t, err, ErrBatchImageGroupDisabled)
- })
- }
- func TestBatchImagePublicService_StatusItemsAndCancel(t *testing.T) {
- ctx := context.Background()
- t.Run("status is owner scoped and maps public status", func(t *testing.T) {
- svc, repo, _, _, _ := newTestBatchImagePublicService(true)
- apiKeyID := int64(22)
- accountID := int64(101)
- repo.jobs["imgbatch_status"] = &BatchImageJob{
- BatchID: "imgbatch_status",
- UserID: 11,
- APIKeyID: &apiKeyID,
- AccountID: &accountID,
- Provider: BatchImageProviderGeminiAPI,
- Model: "gemini-2.5-flash-image",
- Status: BatchImageJobStatusIndexing,
- ProviderJobName: batchImageStringPtr("providers/internal/job"),
- CreatedAt: time.Now(),
- }
- got, err := svc.Get(ctx, testBatchImageOwner(), "imgbatch_status")
- require.NoError(t, err)
- require.Equal(t, "processing_results", got.Status)
- body, err := json.Marshal(got)
- require.NoError(t, err)
- requireBatchImagePublicJSONHasNoInternals(t, string(body))
- _, err = svc.Get(ctx, BatchImageOwner{UserID: 11, APIKeyID: 999}, "imgbatch_status")
- require.ErrorIs(t, err, ErrBatchImageJobNotFound)
- })
- t.Run("items are filtered paginated and sanitized", func(t *testing.T) {
- svc, repo, _, _, _ := newTestBatchImagePublicService(true)
- apiKeyID := int64(22)
- repo.jobs["imgbatch_items"] = &BatchImageJob{
- BatchID: "imgbatch_items",
- UserID: 11,
- APIKeyID: &apiKeyID,
- Provider: BatchImageProviderGeminiAPI,
- Model: "gemini-2.5-flash-image",
- Status: BatchImageJobStatusCompleted,
- CreatedAt: time.Now(),
- }
- sourceObject := "gs://bucket/internal/output.jsonl"
- mime := "image/png"
- ext := "png"
- code := "SAFETY_BLOCKED"
- msg := "blocked in gs://bucket/internal/output.jsonl"
- repo.items["imgbatch_items"] = []CreateBatchImageItemParams{
- {JobID: "imgbatch_items", CustomID: "ok_1", Status: BatchImageItemStatusSuccess, ProviderSourceObject: &sourceObject, MimeType: &mime, FileExtension: &ext, ImageCount: 1},
- {JobID: "imgbatch_items", CustomID: "bad_1", Status: BatchImageItemStatusFailed, ProviderSourceObject: &sourceObject, ErrorCode: &code, ErrorMessage: &msg},
- {JobID: "imgbatch_items", CustomID: "ok_2", Status: BatchImageItemStatusSuccess, MimeType: &mime, FileExtension: &ext, ImageCount: 1},
- }
- page, err := svc.ListItems(ctx, testBatchImageOwner(), "imgbatch_items", BatchImageItemsQuery{Limit: 1})
- require.NoError(t, err)
- require.True(t, page.HasMore)
- require.Len(t, page.Data, 1)
- require.Equal(t, "ok_1", page.Data[0].CustomID)
- filtered, err := svc.ListItems(ctx, testBatchImageOwner(), "imgbatch_items", BatchImageItemsQuery{Status: "failed", Limit: 100})
- require.NoError(t, err)
- require.False(t, filtered.HasMore)
- require.Len(t, filtered.Data, 1)
- require.Equal(t, "failed", filtered.Data[0].Status)
- require.NotNil(t, filtered.Data[0].Error)
- require.Equal(t, "upstream provider operation failed", filtered.Data[0].Error.Message)
- body, err := json.Marshal(filtered)
- require.NoError(t, err)
- requireBatchImagePublicJSONHasNoInternals(t, string(body))
- require.NotContains(t, string(body), "download_url")
- _, err = svc.ListItems(ctx, BatchImageOwner{UserID: 12, APIKeyID: 22}, "imgbatch_items", BatchImageItemsQuery{})
- require.ErrorIs(t, err, ErrBatchImageJobNotFound)
- })
- t.Run("cancel active job calls provider and waits for confirmed terminal state", func(t *testing.T) {
- svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true)
- apiKeyID := int64(22)
- accountID := int64(101)
- holdAmount := 0.5
- holdID := BatchImageHoldRequestID("imgbatch_cancel")
- repo.jobs["imgbatch_cancel"] = &BatchImageJob{
- BatchID: "imgbatch_cancel",
- UserID: 11,
- APIKeyID: &apiKeyID,
- AccountID: &accountID,
- Provider: BatchImageProviderGeminiAPI,
- Model: "gemini-2.5-flash-image",
- Status: BatchImageJobStatusSubmitted,
- ProviderJobName: batchImageStringPtr("providers/internal/job"),
- EstimatedCost: holdAmount,
- HoldAmount: &holdAmount,
- HoldID: &holdID,
- CreatedAt: time.Now(),
- }
- got, err := svc.Cancel(ctx, testBatchImageOwner(), "imgbatch_cancel")
- require.NoError(t, err)
- require.Equal(t, "queued", got.Status)
- require.Equal(t, 1, gemini.cancelCount)
- billing := svc.BillingRepo.(*fakeBatchImageBillingRepo)
- require.Empty(t, billing.releases)
- require.Equal(t, []string{"imgbatch_cancel"}, queue.enqueued)
- require.Equal(t, BatchImageJobStatusSubmitted, repo.jobs["imgbatch_cancel"].Status)
- require.Contains(t, repo.events["imgbatch_cancel"], "job_cancel_requested")
- })
- t.Run("cancel terminal job is idempotent", func(t *testing.T) {
- svc, repo, _, gemini, _ := newTestBatchImagePublicService(true)
- apiKeyID := int64(22)
- repo.jobs["imgbatch_done"] = &BatchImageJob{
- BatchID: "imgbatch_done",
- UserID: 11,
- APIKeyID: &apiKeyID,
- Provider: BatchImageProviderGeminiAPI,
- Model: "gemini-2.5-flash-image",
- Status: BatchImageJobStatusCompleted,
- CreatedAt: time.Now(),
- }
- got, err := svc.Cancel(ctx, testBatchImageOwner(), "imgbatch_done")
- require.NoError(t, err)
- require.Equal(t, "completed", got.Status)
- require.Zero(t, gemini.cancelCount)
- })
- t.Run("cancel hides provider raw errors behind public error", func(t *testing.T) {
- svc, repo, _, gemini, _ := newTestBatchImagePublicService(true)
- gemini.cancelErr = errors.New("projects/secret-provider-job not found")
- apiKeyID := int64(22)
- accountID := int64(101)
- repo.jobs["imgbatch_cancel_error"] = &BatchImageJob{
- BatchID: "imgbatch_cancel_error",
- UserID: 11,
- APIKeyID: &apiKeyID,
- AccountID: &accountID,
- Provider: BatchImageProviderGeminiAPI,
- Model: "gemini-2.5-flash-image",
- Status: BatchImageJobStatusSubmitted,
- ProviderJobName: batchImageStringPtr("providers/internal/job"),
- CreatedAt: time.Now(),
- }
- _, err := svc.Cancel(ctx, testBatchImageOwner(), "imgbatch_cancel_error")
- require.ErrorIs(t, err, ErrBatchImageCancelFailed)
- require.Equal(t, "BATCH_IMAGE_CANCEL_FAILED", infraerrors.Reason(err))
- require.NotContains(t, infraerrors.Message(err), "projects/")
- })
- }
- func newTestBatchImagePublicService(enabled bool) (*BatchImagePublicService, *fakeBatchImageRepository, *publicBatchImageQueue, *publicBatchImageProvider, *publicBatchImageProvider) {
- repo := newFakeBatchImageRepository()
- queue := &publicBatchImageQueue{}
- gemini := &publicBatchImageProvider{name: BatchImageProviderGeminiAPI}
- vertex := &publicBatchImageProvider{name: BatchImageProviderVertex}
- svc := &BatchImagePublicService{
- Repo: repo,
- AccountRepo: &publicBatchImageAccountRepo{accounts: []Account{testBatchImageAccount(101, AccountTypeAPIKey), testBatchImageAccount(202, AccountTypeServiceAccount)}},
- Queue: queue,
- ProviderRegistry: NewBatchImageProviderRegistry(
- gemini,
- vertex,
- ),
- Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25},
- BillingRepo: &fakeBatchImageBillingRepo{},
- AuthCache: &fakeBatchImageAuthCacheInvalidator{},
- Config: &config.Config{BatchImage: config.BatchImageConfig{
- Enabled: enabled,
- MaxItemsPerJobDefault: 2,
- MaxPromptCharsPerItem: 8,
- DefaultResponseMimeType: "image/png",
- DefaultImageSize: "1K",
- }},
- }
- return svc, repo, queue, gemini, vertex
- }
- func testBatchImageOwner() BatchImageOwner {
- return BatchImageOwner{UserID: 11, APIKeyID: 22}
- }
- type fakeBatchImageAuthCacheInvalidator struct {
- keys []string
- userIDs []int64
- groupIDs []int64
- }
- func (f *fakeBatchImageAuthCacheInvalidator) InvalidateAuthCacheByKey(_ context.Context, key string) {
- f.keys = append(f.keys, key)
- }
- func (f *fakeBatchImageAuthCacheInvalidator) InvalidateAuthCacheByUserID(_ context.Context, userID int64) {
- f.userIDs = append(f.userIDs, userID)
- }
- func (f *fakeBatchImageAuthCacheInvalidator) InvalidateAuthCacheByGroupID(_ context.Context, groupID int64) {
- f.groupIDs = append(f.groupIDs, groupID)
- }
- func validBatchImageSubmitRequest() BatchImageSubmitRequest {
- return BatchImageSubmitRequest{
- Model: "gemini-2.5-flash-image",
- Provider: BatchImageProviderGeminiAPI,
- ResponseMimeType: "image/png",
- AspectRatio: "1:1",
- ImageSize: "1K",
- Metadata: map[string]string{"project": "campaign-a", "secret": strings.Repeat("x", 300)},
- Items: []BatchImageSubmitItem{
- {CustomID: "cover_001", Prompt: "hero"},
- {CustomID: "cover_002", Prompt: "clean"},
- },
- }
- }
- func testBatchImageAccount(id int64, accountType string) Account {
- return Account{
- ID: id,
- Platform: PlatformGemini,
- Type: accountType,
- Status: StatusActive,
- Schedulable: true,
- Priority: int(id),
- Credentials: map[string]any{"api_key": "test-secret"},
- Concurrency: 1,
- RateLimitedAt: nil,
- }
- }
- func testBatchImageMappedAccount(id int64, accountType string, mapping map[string]any) Account {
- account := testBatchImageAccount(id, accountType)
- account.Credentials["model_mapping"] = mapping
- return account
- }
- func requireBatchImagePublicJSONHasNoInternals(t *testing.T, body string) {
- t.Helper()
- for _, forbidden := range []string{
- "provider_job_name",
- "provider_input_ref",
- "provider_output_ref",
- "gcs_input_uri",
- "gcs_output_uri",
- "account_id",
- "service_account",
- "api_key",
- "download_url",
- "providers/",
- "files/",
- "gs://",
- } {
- require.NotContains(t, body, forbidden)
- }
- }
- type publicBatchImageAccountRepo struct {
- accounts []Account
- }
- func (r *publicBatchImageAccountRepo) GetByID(_ context.Context, id int64) (*Account, error) {
- for i := range r.accounts {
- if r.accounts[i].ID == id {
- return &r.accounts[i], nil
- }
- }
- return nil, errors.New("account not found")
- }
- func (r *publicBatchImageAccountRepo) ListSchedulableByPlatform(_ context.Context, platform string) ([]Account, error) {
- out := make([]Account, 0, len(r.accounts))
- for _, account := range r.accounts {
- if account.Platform == platform {
- out = append(out, account)
- }
- }
- return out, nil
- }
- func (r *publicBatchImageAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, _ int64, platform string) ([]Account, error) {
- return r.ListSchedulableByPlatform(ctx, platform)
- }
- type publicBatchImageQueue struct {
- enqueued []string
- err error
- }
- func (q *publicBatchImageQueue) Enqueue(_ context.Context, batchID string) error {
- if q.err != nil {
- return q.err
- }
- for _, existing := range q.enqueued {
- if existing == batchID {
- return ErrBatchImageAlreadyQueued
- }
- }
- q.enqueued = append(q.enqueued, batchID)
- return nil
- }
- func (q *publicBatchImageQueue) Reserve(context.Context, time.Duration) (ReservedBatchImageJob, error) {
- return ReservedBatchImageJob{}, ErrBatchImageQueueEmpty
- }
- func (q *publicBatchImageQueue) RequeueAfter(context.Context, string, time.Duration) error {
- return nil
- }
- func (q *publicBatchImageQueue) Ack(context.Context, string) error {
- return nil
- }
- func (q *publicBatchImageQueue) Heartbeat(context.Context, string) error {
- return nil
- }
- func (q *publicBatchImageQueue) MoveDueDelayedToReady(context.Context, int) (int, error) {
- return 0, nil
- }
- func (q *publicBatchImageQueue) RecoverStaleActive(context.Context, time.Duration, int) (int, error) {
- return 0, nil
- }
- func (q *publicBatchImageQueue) TryAcquireJobLock(context.Context, string, time.Duration) (BatchImageJobLock, bool, error) {
- return nil, false, nil
- }
- type publicBatchImageProvider struct {
- name string
- submits []BatchImageInput
- submitErr error
- cancelCount int
- cancelErr error
- result string
- cleanupTargets []CleanupTarget
- cleanupErr error
- }
- func (p *publicBatchImageProvider) Name() string { return p.name }
- func (p *publicBatchImageProvider) SupportsAccount(*Account) bool { return true }
- func (p *publicBatchImageProvider) Submit(_ context.Context, _ *BatchImageJob, _ *Account, input BatchImageInput) (*BatchProviderJob, error) {
- p.submits = append(p.submits, input)
- if p.submitErr != nil {
- return nil, p.submitErr
- }
- return &BatchProviderJob{
- ProviderJobName: "providers/" + p.name + "/job",
- ProviderInputRef: "files/" + p.name + "/input",
- ProviderOutputRef: "files/" + p.name + "/output",
- }, nil
- }
- func (p *publicBatchImageProvider) Get(context.Context, *BatchImageJob, *Account) (*BatchProviderStatus, error) {
- return &BatchProviderStatus{InternalState: BatchProviderStateQueued}, nil
- }
- func (p *publicBatchImageProvider) Cancel(context.Context, *BatchImageJob, *Account) error {
- p.cancelCount++
- return p.cancelErr
- }
- func (p *publicBatchImageProvider) OpenResult(context.Context, *BatchImageJob, *Account) (io.ReadCloser, string, error) {
- return io.NopCloser(strings.NewReader(p.result)), "application/jsonl", nil
- }
- func (p *publicBatchImageProvider) Cleanup(_ context.Context, _ *BatchImageJob, _ *Account, target CleanupTarget) error {
- p.cleanupTargets = append(p.cleanupTargets, target)
- return p.cleanupErr
- }
- var _ BatchImageAccountSelectionRepository = (*publicBatchImageAccountRepo)(nil)
- var _ BatchImageQueue = (*publicBatchImageQueue)(nil)
- var _ BatchImageProvider = (*publicBatchImageProvider)(nil)
- type publicBatchImageGroupRepo struct {
- groups map[int64]*Group
- }
- func (r *publicBatchImageGroupRepo) GetByIDLite(_ context.Context, id int64) (*Group, error) {
- if r != nil && r.groups != nil {
- if group, ok := r.groups[id]; ok {
- return group, nil
- }
- }
- return nil, ErrGroupNotFound
- }
- type publicBatchImageUserGroupRateRepo struct {
- rates map[int64]*float64
- }
- func (r *publicBatchImageUserGroupRateRepo) GetByUserAndGroup(_ context.Context, _ int64, groupID int64) (*float64, error) {
- if r != nil && r.rates != nil {
- return r.rates[groupID], nil
- }
- return nil, nil
- }
- var _ BatchImageGroupPricingRepository = (*publicBatchImageGroupRepo)(nil)
- var _ BatchImageUserGroupRateRepository = (*publicBatchImageUserGroupRateRepo)(nil)
|