account_test_service_gemini_test.go 1.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859
  1. //go:build unit
  2. package service
  3. import (
  4. "encoding/json"
  5. "strings"
  6. "testing"
  7. "github.com/gin-gonic/gin"
  8. "github.com/stretchr/testify/require"
  9. )
  10. func TestCreateGeminiTestPayload_ImageModel(t *testing.T) {
  11. t.Parallel()
  12. payload := createGeminiTestPayload("gemini-2.5-flash-image", "draw a tiny robot")
  13. var parsed struct {
  14. Contents []struct {
  15. Parts []struct {
  16. Text string `json:"text"`
  17. } `json:"parts"`
  18. } `json:"contents"`
  19. GenerationConfig struct {
  20. ResponseModalities []string `json:"responseModalities"`
  21. ImageConfig struct {
  22. AspectRatio string `json:"aspectRatio"`
  23. } `json:"imageConfig"`
  24. } `json:"generationConfig"`
  25. }
  26. require.NoError(t, json.Unmarshal(payload, &parsed))
  27. require.Len(t, parsed.Contents, 1)
  28. require.Len(t, parsed.Contents[0].Parts, 1)
  29. require.Equal(t, "draw a tiny robot", parsed.Contents[0].Parts[0].Text)
  30. require.Equal(t, []string{"TEXT", "IMAGE"}, parsed.GenerationConfig.ResponseModalities)
  31. require.Equal(t, "1:1", parsed.GenerationConfig.ImageConfig.AspectRatio)
  32. }
  33. func TestProcessGeminiStream_EmitsImageEvent(t *testing.T) {
  34. t.Parallel()
  35. gin.SetMode(gin.TestMode)
  36. ctx, recorder := newTestContext()
  37. svc := &AccountTestService{}
  38. stream := strings.NewReader("data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"ok\"},{\"inlineData\":{\"mimeType\":\"image/png\",\"data\":\"QUJD\"}}]}}]}\n\ndata: [DONE]\n\n")
  39. err := svc.processGeminiStream(ctx, stream)
  40. require.NoError(t, err)
  41. body := recorder.Body.String()
  42. require.Contains(t, body, "\"type\":\"content\"")
  43. require.Contains(t, body, "\"text\":\"ok\"")
  44. require.Contains(t, body, "\"type\":\"image\"")
  45. require.Contains(t, body, "\"image_url\":\"data:image/png;base64,QUJD\"")
  46. require.Contains(t, body, "\"mime_type\":\"image/png\"")
  47. }