301 lines
9.6 KiB
Go
301 lines
9.6 KiB
Go
package llmtests
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/base64"
|
|
"fmt"
|
|
"image"
|
|
_ "image/jpeg"
|
|
_ "image/png"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
bifrost "github.com/maximhq/bifrost/core"
|
|
"github.com/maximhq/bifrost/core/schemas"
|
|
)
|
|
|
|
// RunImageGenerationTest executes the end-to-end image generation test (non-streaming)
|
|
func RunImageGenerationTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
|
|
if !testConfig.Scenarios.ImageGeneration {
|
|
t.Logf("Image generation not supported for provider %s", testConfig.Provider)
|
|
return
|
|
}
|
|
|
|
if testConfig.ImageGenerationModel == "" {
|
|
t.Logf("Image generation not configured for provider %s", testConfig.Provider)
|
|
return
|
|
}
|
|
|
|
t.Run("ImageGeneration", func(t *testing.T) {
|
|
if os.Getenv("SKIP_PARALLEL_TESTS") != "true" {
|
|
t.Parallel()
|
|
}
|
|
|
|
retryConfig := GetTestRetryConfigForScenario("ImageGeneration", testConfig)
|
|
retryContext := TestRetryContext{
|
|
ScenarioName: "ImageGeneration",
|
|
ExpectedBehavior: map[string]interface{}{},
|
|
TestMetadata: map[string]interface{}{
|
|
"provider": testConfig.Provider,
|
|
"model": testConfig.ImageGenerationModel,
|
|
},
|
|
}
|
|
|
|
expectations := GetExpectationsForScenario("ImageGeneration", testConfig, map[string]interface{}{
|
|
"min_images": 1,
|
|
"expected_size": "1024x1024",
|
|
})
|
|
|
|
imageGenerationRetryConfig := ImageGenerationRetryConfig{
|
|
MaxAttempts: retryConfig.MaxAttempts,
|
|
BaseDelay: retryConfig.BaseDelay,
|
|
MaxDelay: retryConfig.MaxDelay,
|
|
Conditions: []ImageGenerationRetryCondition{},
|
|
OnRetry: retryConfig.OnRetry,
|
|
OnFinalFail: retryConfig.OnFinalFail,
|
|
}
|
|
// Test basic image generation
|
|
imageGenerationOperation := func() (*schemas.BifrostImageGenerationResponse, *schemas.BifrostError) {
|
|
request := &schemas.BifrostImageGenerationRequest{
|
|
Provider: testConfig.Provider,
|
|
Model: testConfig.ImageGenerationModel,
|
|
Input: &schemas.ImageGenerationInput{
|
|
Prompt: "A serene Japanese garden with cherry blossoms in spring",
|
|
},
|
|
Params: &schemas.ImageGenerationParameters{
|
|
Size: bifrost.Ptr("1024x1024"),
|
|
N: bifrost.Ptr(1),
|
|
},
|
|
Fallbacks: testConfig.ImageGenerationFallbacks,
|
|
}
|
|
|
|
response, err := client.ImageGenerationRequest(schemas.NewBifrostContext(ctx, schemas.NoDeadline), request)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if response != nil {
|
|
return response, nil
|
|
}
|
|
return nil, &schemas.BifrostError{
|
|
IsBifrostError: true,
|
|
Error: &schemas.ErrorField{
|
|
Message: "No image generation response returned",
|
|
},
|
|
}
|
|
}
|
|
|
|
imageGenerationResponse, imageGenerationError := WithImageGenerationRetry(t, imageGenerationRetryConfig, retryContext, expectations, "ImageGeneration", imageGenerationOperation)
|
|
|
|
if imageGenerationError != nil {
|
|
t.Fatalf("❌ Image generation failed: %v", GetErrorMessage(imageGenerationError))
|
|
}
|
|
|
|
// Validate response
|
|
if imageGenerationResponse == nil {
|
|
t.Fatal("❌ Image generation returned nil response")
|
|
}
|
|
|
|
if len(imageGenerationResponse.Data) == 0 {
|
|
t.Fatal("❌ Image generation returned no image data")
|
|
}
|
|
|
|
// Validate first image
|
|
imageData := imageGenerationResponse.Data[0]
|
|
if imageData.B64JSON == "" && imageData.URL == "" {
|
|
t.Fatal("❌ Image data missing both b64_json and URL")
|
|
}
|
|
|
|
// Validate base64 if present
|
|
if imageData.B64JSON != "" {
|
|
// Decode base64 image data
|
|
decoded, err := base64.StdEncoding.DecodeString(imageData.B64JSON)
|
|
if err != nil {
|
|
t.Fatalf("❌ Failed to decode base64 image data: %v", err)
|
|
}
|
|
if len(decoded) == 0 {
|
|
t.Fatalf("❌ Decoded image data is empty")
|
|
}
|
|
|
|
// Decode image config to validate dimensions
|
|
reader := bytes.NewReader(decoded)
|
|
config, format, err := image.DecodeConfig(reader)
|
|
if err != nil {
|
|
t.Fatalf("❌ Failed to decode image config: %v (format: %s)", err, format)
|
|
}
|
|
|
|
// Validate dimensions are 1024x1024 as requested
|
|
expectedWidth, expectedHeight := 1024, 1024
|
|
if config.Width != expectedWidth || config.Height != expectedHeight {
|
|
t.Errorf("❌ Image dimensions mismatch: got %dx%d, expected %dx%d", config.Width, config.Height, expectedWidth, expectedHeight)
|
|
}
|
|
}
|
|
|
|
// Validate usage if present
|
|
if imageGenerationResponse.Usage != nil {
|
|
if imageGenerationResponse.Usage.TotalTokens == 0 {
|
|
t.Logf("⚠️ Usage total_tokens is 0 (may be provider-specific)")
|
|
}
|
|
}
|
|
|
|
// Validate extra fields
|
|
if imageGenerationResponse.ExtraFields.Provider == "" {
|
|
t.Error("❌ ExtraFields.Provider is empty")
|
|
}
|
|
|
|
if imageGenerationResponse.ExtraFields.OriginalModelRequested == "" {
|
|
t.Error("❌ ExtraFields.OriginalModelRequested is empty")
|
|
}
|
|
|
|
t.Logf("✅ Image generation successful: ID=%s, Provider=%s, Model=%s, Images=%d",
|
|
imageGenerationResponse.ID, imageGenerationResponse.ExtraFields.Provider, imageGenerationResponse.ExtraFields.OriginalModelRequested, len(imageGenerationResponse.Data))
|
|
})
|
|
}
|
|
|
|
// RunImageGenerationStreamTest executes the end-to-end streaming image generation test
|
|
func RunImageGenerationStreamTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
|
|
if !testConfig.Scenarios.ImageGenerationStream {
|
|
t.Logf("Image generation streaming not supported for provider %s", testConfig.Provider)
|
|
return
|
|
}
|
|
|
|
if testConfig.ImageGenerationModel == "" {
|
|
t.Logf("Image generation streaming not configured for provider %s", testConfig.Provider)
|
|
return
|
|
}
|
|
|
|
t.Run("ImageGenerationStream", func(t *testing.T) {
|
|
if os.Getenv("SKIP_PARALLEL_TESTS") != "true" {
|
|
t.Parallel()
|
|
}
|
|
|
|
retryConfig := GetTestRetryConfigForScenario("ImageGenerationStream", testConfig)
|
|
retryContext := TestRetryContext{
|
|
ScenarioName: "ImageGenerationStream",
|
|
ExpectedBehavior: map[string]interface{}{
|
|
"should_generate_images": true,
|
|
},
|
|
TestMetadata: map[string]interface{}{
|
|
"provider": testConfig.Provider,
|
|
"model": testConfig.ImageGenerationModel,
|
|
},
|
|
}
|
|
|
|
request := &schemas.BifrostImageGenerationRequest{
|
|
Provider: testConfig.Provider,
|
|
Model: testConfig.ImageGenerationModel,
|
|
Input: &schemas.ImageGenerationInput{
|
|
Prompt: "A futuristic cityscape at sunset with flying cars",
|
|
},
|
|
Params: &schemas.ImageGenerationParameters{
|
|
Size: bifrost.Ptr("1024x1024"),
|
|
Quality: bifrost.Ptr("low"),
|
|
},
|
|
Fallbacks: testConfig.ImageGenerationFallbacks,
|
|
}
|
|
streamCtx, cancel := context.WithTimeout(ctx, 60*time.Second)
|
|
defer cancel()
|
|
|
|
validationResult := WithImageGenerationStreamRetry(
|
|
t,
|
|
retryConfig,
|
|
retryContext,
|
|
func() (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) {
|
|
return client.ImageGenerationStreamRequest(schemas.NewBifrostContext(streamCtx, schemas.NoDeadline), request)
|
|
},
|
|
func(responseChannel chan *schemas.BifrostStreamChunk) ImageGenerationStreamValidationResult {
|
|
// Validate stream content
|
|
var receivedData bool
|
|
var streamErrors []string
|
|
var validationErrors []string
|
|
hasCompleted := false
|
|
|
|
for {
|
|
select {
|
|
case response, ok := <-responseChannel:
|
|
if !ok {
|
|
goto streamComplete
|
|
}
|
|
|
|
if response == nil {
|
|
streamErrors = append(streamErrors, "Received nil stream response")
|
|
continue
|
|
}
|
|
|
|
if response.BifrostError != nil {
|
|
streamErrors = append(streamErrors, fmt.Sprintf("Error in stream: %s", GetErrorMessage(response.BifrostError)))
|
|
continue
|
|
}
|
|
|
|
if response.BifrostImageGenerationStreamResponse != nil {
|
|
receivedData = true
|
|
imgResp := response.BifrostImageGenerationStreamResponse
|
|
|
|
if imgResp.Type == schemas.ImageGenerationEventTypeCompleted {
|
|
hasCompleted = true
|
|
// Validate that completed images have actual data
|
|
if imgResp.URL == "" && imgResp.B64JSON == "" {
|
|
validationErrors = append(validationErrors, "Completion chunk received but image has no URL or B64JSON data")
|
|
}
|
|
}
|
|
}
|
|
case <-streamCtx.Done():
|
|
validationErrors = append(validationErrors, "Stream validation timed out")
|
|
drainCtx, drainCancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
go func() {
|
|
defer drainCancel()
|
|
for {
|
|
select {
|
|
case _, ok := <-responseChannel:
|
|
if !ok {
|
|
return
|
|
}
|
|
case <-drainCtx.Done():
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
goto streamComplete
|
|
}
|
|
}
|
|
streamComplete:
|
|
|
|
// Stream errors should cause the test to fail - convert them to validation errors
|
|
if len(streamErrors) > 0 {
|
|
validationErrors = append(validationErrors, fmt.Sprintf("Stream errors encountered: %s", strings.Join(streamErrors, "; ")))
|
|
}
|
|
|
|
// Test passes only if: data received, completion received, and no errors (including stream errors)
|
|
passed := receivedData && hasCompleted && len(validationErrors) == 0
|
|
if !receivedData {
|
|
validationErrors = append(validationErrors, "No stream data received")
|
|
}
|
|
if !hasCompleted {
|
|
validationErrors = append(validationErrors, "No completion chunk received")
|
|
}
|
|
|
|
return ImageGenerationStreamValidationResult{
|
|
Passed: passed,
|
|
Errors: validationErrors,
|
|
ReceivedData: receivedData,
|
|
StreamErrors: streamErrors,
|
|
}
|
|
},
|
|
)
|
|
|
|
if !validationResult.Passed {
|
|
allErrors := append(validationResult.Errors, validationResult.StreamErrors...)
|
|
t.Fatalf("❌ Image generation stream validation failed: %s", strings.Join(allErrors, "; "))
|
|
}
|
|
|
|
if !validationResult.ReceivedData {
|
|
t.Fatal("❌ No stream data received")
|
|
}
|
|
|
|
t.Logf("✅ Image generation stream successful: ReceivedData=%v, Errors=%d, StreamErrors=%d",
|
|
validationResult.ReceivedData, len(validationResult.Errors), len(validationResult.StreamErrors))
|
|
})
|
|
}
|