azure and bedrock test case fixes

Author: akshaydeo

.github/workflows/scripts/release-core.sh

@@ -55,7 +55,7 @@ echo "✅ MCP test servers built"
 # Run core tests with coverage
 echo "🔧 Running core tests with coverage..."
 cd core
-go test -race -v -timeout 20m -coverprofile=coverage.txt -coverpkg=./... ./...
+go test -race -timeout 20m -coverprofile=coverage.txt -coverpkg=./... ./...
 
 # Upload coverage to Codecov
 if [ -n "${CODECOV_TOKEN:-}" ]; then

core/internal/llmtests/batch.go

@@ -18,11 +18,11 @@ func getFakeBatchID(provider schemas.ModelProvider) string {
 	case schemas.OpenAI, schemas.Azure:
 		return "batch_test-batch-id"
 	case schemas.Bedrock:
-		// Bedrock uses ARNs for batch IDs
-		return "arn:aws:bedrock:us-east-1:791152688819:model-invocation-job/test-batch-id"
+		// Bedrock uses ARNs for batch IDs - job ID must be exactly 12 lowercase alphanumeric chars
+		return "arn:aws:bedrock:us-east-1:791152688819:model-invocation-job/aaaaaaaaaaaa"
 	case schemas.Gemini:
-		// Gemini uses "batches/" prefix for batch IDs
-		return "batches/test-batch-id"
+		// Gemini uses "batches/" prefix with alphanumeric batch IDs
+		return "batches/aaaaaaaaaaaa"
 	default:
 		return "batch_test-batch-id"
 	}
@@ -464,23 +464,30 @@ func RunBatchResultsTest(t *testing.T, client *bifrost.Bifrost, ctx context.Cont
 				}
 			}
 
-			// Check error type/message for "not found" indicators
-			isNotFoundError := false
+			// Check error type/message for "not found" or validation error indicators
+			// Providers may return "not found" errors OR validation errors for fake batch IDs
+			isExpectedError := false
 			if err.Error != nil {
 				if err.Error.Type != nil {
 					errType := strings.ToLower(*err.Error.Type)
 					if strings.Contains(errType, "not_found") || strings.Contains(errType, "notfound") {
-						isNotFoundError = true
+						isExpectedError = true
 					}
 				}
 				errMsg := strings.ToLower(err.Error.Message)
+				// Check for "not found" style errors
 				if strings.Contains(errMsg, "not found") || strings.Contains(errMsg, "does not exist") || strings.Contains(errMsg, "could not find") {
-					isNotFoundError = true
+					isExpectedError = true
+				}
+				// Check for validation/parsing errors (providers validate batch ID format before lookup)
+				if strings.Contains(errMsg, "could not parse") || strings.Contains(errMsg, "validation error") ||
+					strings.Contains(errMsg, "failed to satisfy constraint") || strings.Contains(errMsg, "invalid") {
+					isExpectedError = true
 				}
 			}
 
-			if isNotFoundError {
-				t.Logf("[INFO] BatchResults test completed (expected not found error with fake ID): %v", GetErrorMessage(err))
+			if isExpectedError {
+				t.Logf("[INFO] BatchResults test completed (expected error with fake ID): %v", GetErrorMessage(err))
 				return
 			}
 

core/internal/llmtests/image_generation.go

@@ -19,6 +19,11 @@ import (
 
 // 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

core/providers/azure/azure.go

@@ -1105,41 +1105,50 @@ func (provider *AzureProvider) SpeechStream(ctx *schemas.BifrostContext, postHoo
 						audioData = event
 					}
 
-					// First, try to parse as JSON error response (these would be valid JSON text)
-					var bifrostErr schemas.BifrostError
-					if err := sonic.Unmarshal(audioData, &bifrostErr); err == nil {
-						if bifrostErr.Error != nil && bifrostErr.Error.Message != "" {
-							bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
-								Provider:       provider.GetProviderKey(),
-								ModelRequested: request.Model,
-								RequestType:    schemas.SpeechStreamRequest,
+					// Skip empty data
+					if len(audioData) == 0 {
+						continue
+					}
+
+					// Azure sends JSON-wrapped responses for speech streaming
+					// Parse the JSON to extract the response type and audio data
+					var response schemas.BifrostSpeechStreamResponse
+					if err := sonic.Unmarshal(audioData, &response); err != nil {
+						// If JSON parsing fails, check if this might be an error response
+						var bifrostErr schemas.BifrostError
+						if errParseErr := sonic.Unmarshal(audioData, &bifrostErr); errParseErr == nil {
+							if bifrostErr.Error != nil && bifrostErr.Error.Message != "" {
+								bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
+									Provider:       provider.GetProviderKey(),
+									ModelRequested: request.Model,
+									RequestType:    schemas.SpeechStreamRequest,
+								}
+								ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true)
+								providerUtils.ProcessAndSendBifrostError(ctx, postHookRunner, &bifrostErr, responseChan, provider.logger)
+								return
 							}
-							ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true)
-							providerUtils.ProcessAndSendBifrostError(ctx, postHookRunner, &bifrostErr, responseChan, provider.logger)
-							return
 						}
+						// If it's not valid JSON, log and skip
+						provider.logger.Warn("failed to parse speech stream response: %v", err)
+						continue
 					}
 
-					// Skip empty audio data
-					if len(audioData) == 0 {
+					// Check for completion event - skip if no audio data
+					if response.Type == schemas.SpeechStreamResponseTypeDone || len(response.Audio) == 0 {
+						// This is a control event or empty response - skip
 						continue
 					}
 
 					chunkIndex++
 
-					// Create response with raw audio data
-					// Azure sends raw binary MP3 frames (starting with 0xff 0xf3 or 0xff 0xfb)
-					response := schemas.BifrostSpeechStreamResponse{
-						Type:  schemas.SpeechStreamResponseTypeDelta,
-						Audio: audioData,
-						ExtraFields: schemas.BifrostResponseExtraFields{
-							RequestType:     schemas.SpeechStreamRequest,
-							Provider:        provider.GetProviderKey(),
-							ModelRequested:  request.Model,
-							ModelDeployment: deployment,
-							ChunkIndex:      chunkIndex,
-							Latency:         time.Since(lastChunkTime).Milliseconds(),
-						},
+					// Set extra fields for the response
+					response.ExtraFields = schemas.BifrostResponseExtraFields{
+						RequestType:     schemas.SpeechStreamRequest,
+						Provider:        provider.GetProviderKey(),
+						ModelRequested:  request.Model,
+						ModelDeployment: deployment,
+						ChunkIndex:      chunkIndex,
+						Latency:         time.Since(lastChunkTime).Milliseconds(),
 					}
 					lastChunkTime = time.Now()
 

core/providers/bedrock/chat.go

@@ -43,10 +43,6 @@ func ToBedrockChatCompletionRequest(ctx *schemas.BifrostContext, bifrostReq *sch
 	// Ensure tool config is present when needed
 	ensureChatToolConfigForConversation(bifrostReq, bedrockReq)
 
-	if bifrostReq.Params != nil {
-		bedrockReq.ExtraParams = bifrostReq.Params.ExtraParams
-	}
-
 	return bedrockReq, nil
 }
 

core/providers/bedrock/utils.go

@@ -207,18 +207,15 @@ func convertChatParameters(ctx *schemas.BifrostContext, bifrostReq *schemas.Bifr
 				config := &BedrockGuardrailConfig{}
 
 				if identifier, ok := gc["guardrailIdentifier"].(string); ok {
-					delete(gc, "guardrailIdentifier")
 					config.GuardrailIdentifier = identifier
 				}
 				if version, ok := gc["guardrailVersion"].(string); ok {
-					delete(gc, "guardrailVersion")
 					config.GuardrailVersion = version
 				}
 				if trace, ok := gc["trace"].(string); ok {
-					delete(gc, "trace")
 					config.Trace = &trace
 				}
-				bedrockReq.ExtraParams["guardrailConfig"] = gc
+				delete(bedrockReq.ExtraParams, "guardrailConfig")
 				bedrockReq.GuardrailConfig = config
 			}
 		}
@@ -253,13 +250,11 @@ func convertChatParameters(ctx *schemas.BifrostContext, bifrostReq *schemas.Bifr
 			// Handle performance configuration
 			if perfConfig, exists := bifrostReq.Params.ExtraParams["performanceConfig"]; exists {
 				if pc, ok := perfConfig.(map[string]interface{}); ok {
-
 					config := &BedrockPerformanceConfig{}
 					if latency, ok := pc["latency"].(string); ok {
-						delete(pc, "latency")
 						config.Latency = &latency
 					}
-					bedrockReq.ExtraParams["performanceConfig"] = pc
+					delete(bedrockReq.ExtraParams, "performanceConfig")
 					bedrockReq.PerformanceConfig = config
 				}
 			}
@@ -292,6 +287,10 @@ func convertChatParameters(ctx *schemas.BifrostContext, bifrostReq *schemas.Bifr
 				}
 			}
 		}
+		// Set ExtraParams to nil if all keys were extracted to dedicated fields
+		if len(bedrockReq.ExtraParams) == 0 {
+			bedrockReq.ExtraParams = nil
+		}
 	}
 	return nil
 }