test: improve OpenInference image generation tests

Signed-off-by: Hrushikesh Patil <[email protected]>

Author: nutanix-Hrushikesh

internal/tracing/openinference/openai/image_generation_test.go

@@ -6,80 +6,167 @@
 package openai
 
 import (
+	"encoding/json"
 	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/codes"
 	"go.opentelemetry.io/otel/sdk/trace"
-	"go.opentelemetry.io/otel/sdk/trace/tracetest"
+	oteltrace "go.opentelemetry.io/otel/trace"
 
+	"github.com/envoyproxy/ai-gateway/internal/testing/testotel"
 	tracing "github.com/envoyproxy/ai-gateway/internal/tracing/api"
 	openaisdk "github.com/openai/openai-go/v2"
 	openaiparam "github.com/openai/openai-go/v2/packages/param"
 )
 
-func TestImageGenerationRecorder_RequestAttributes_SDK(t *testing.T) {
-	exporter := tracetest.NewInMemoryExporter()
-	tp := trace.NewTracerProvider(trace.WithSyncer(exporter))
-	tr := tp.Tracer("test")
-
-	recorder := NewImageGenerationRecorder(nil) // default config
-
-	params := &openaisdk.ImageGenerateParams{
+var (
+	// Test data constants following chat completion pattern
+	basicImageReq = &openaisdk.ImageGenerateParams{
 		Model:          openaisdk.ImageModelGPTImage1,
 		Prompt:         "a hummingbird",
 		Size:           openaisdk.ImageGenerateParamsSize1024x1024,
 		Quality:        openaisdk.ImageGenerateParamsQualityHigh,
 		ResponseFormat: openaisdk.ImageGenerateParamsResponseFormatB64JSON,
-		N:              openaiparam.NewOpt[int64](2),
+		N:              openaiparam.NewOpt[int64](1),
 	}
+	basicImageReqBody = mustJSON(basicImageReq)
+
+	basicImageResp = &openaisdk.ImagesResponse{
+		Data: []openaisdk.Image{{URL: "https://example.com/img.png"}},
+		Size: openaisdk.ImagesResponseSize1024x1024,
+		Usage: openaisdk.ImagesResponseUsage{
+			InputTokens:  8,
+			OutputTokens: 1056,
+			TotalTokens:  1064,
+		},
+	}
+	basicImageRespBody = mustJSON(basicImageResp)
+)
 
-	spanName, opts := recorder.StartParams(params, []byte("{}"))
-	require.Equal(t, "ImageGeneration", spanName)
-
-	_, span := tr.Start(t.Context(), spanName, opts...)
-	recorder.RecordRequest(span, params, []byte(`{"prompt":"a hummingbird"}`))
-	span.End()
-
-	spans := exporter.GetSpans()
-	require.Len(t, spans, 1)
-	got := spans[0]
-
-	// Verify a subset of attributes were recorded
-	attrs := attributesToMap(got.Attributes)
-	require.Equal(t, "openai", attrs["llm.system"])                                  // LLMSystemOpenAI
-	require.Equal(t, string(openaisdk.ImageModelGPTImage1), attrs["llm.model_name"]) // model
-	require.Equal(t, "{\"prompt\":\"a hummingbird\"}", attrs["input.value"])         // input body json
-	require.Equal(t, "image_generation", attrs["gen_ai.operation.name"])             // operation name
-	require.Equal(t, "1024x1024", attrs["gen_ai.image.size"])                        // size
-	require.Equal(t, "high", attrs["gen_ai.image.quality"])                          // quality
-	require.Equal(t, "b64_json", attrs["gen_ai.image.response_format"])              // response_format
-	require.Equal(t, "2", attrs["gen_ai.image.n"])                                   // n
+func TestImageGenerationRecorder_StartParams(t *testing.T) {
+	tests := []struct {
+		name             string
+		req              *openaisdk.ImageGenerateParams
+		reqBody          []byte
+		expectedSpanName string
+	}{
+		{
+			name:             "basic request",
+			req:              basicImageReq,
+			reqBody:          basicImageReqBody,
+			expectedSpanName: "ImageGeneration",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			recorder := NewImageGenerationRecorder(nil)
+
+			spanName, opts := recorder.StartParams(tt.req, tt.reqBody)
+			actualSpan := testotel.RecordNewSpan(t, spanName, opts...)
+
+			require.Equal(t, tt.expectedSpanName, actualSpan.Name)
+			require.Equal(t, oteltrace.SpanKindInternal, actualSpan.SpanKind)
+		})
+	}
 }
 
-func TestImageGenerationRecorder_ResponseAttributes_SDK(t *testing.T) {
-	exporter := tracetest.NewInMemoryExporter()
-	tp := trace.NewTracerProvider(trace.WithSyncer(exporter))
-	tr := tp.Tracer("test")
+func TestImageGenerationRecorder_RecordRequest(t *testing.T) {
+	tests := []struct {
+		name          string
+		req           *openaisdk.ImageGenerateParams
+		reqBody       []byte
+		expectedAttrs []attribute.KeyValue
+	}{
+		{
+			name:    "basic request",
+			req:     basicImageReq,
+			reqBody: basicImageReqBody,
+			expectedAttrs: []attribute.KeyValue{
+				attribute.String("gen_ai.operation.name", "image_generation"),
+				attribute.String("gen_ai.image.size", "1024x1024"),
+				attribute.String("gen_ai.image.quality", "high"),
+				attribute.String("gen_ai.image.response_format", "b64_json"),
+				attribute.String("gen_ai.image.n", "1"),
+			},
+		},
+	}
 
-	recorder := NewImageGenerationRecorder(nil)
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			recorder := NewImageGenerationRecorder(nil)
+
+			actualSpan := testotel.RecordWithSpan(t, func(span oteltrace.Span) bool {
+				recorder.RecordRequest(span, tt.req, tt.reqBody)
+				return false
+			})
+
+			// Check that key attributes are present
+			attrs := attributesToMap(actualSpan.Attributes)
+			require.Equal(t, "image_generation", attrs["gen_ai.operation.name"])
+			require.Equal(t, "1024x1024", attrs["gen_ai.image.size"])
+			require.Equal(t, "high", attrs["gen_ai.image.quality"])
+			require.Equal(t, "b64_json", attrs["gen_ai.image.response_format"])
+			require.Equal(t, "1", attrs["gen_ai.image.n"])
+		})
+	}
+}
 
-	_, span := tr.Start(t.Context(), "ImageGeneration")
-	resp := &openaisdk.ImagesResponse{
-		Data: []openaisdk.Image{{URL: "https://example.com/img.png"}},
-		Size: openaisdk.ImagesResponseSize1024x1024,
+func TestImageGenerationRecorder_RecordResponse(t *testing.T) {
+	tests := []struct {
+		name           string
+		respBody       []byte
+		expectedAttrs  []attribute.KeyValue
+		expectedStatus trace.Status
+	}{
+		{
+			name:     "successful response",
+			respBody: basicImageRespBody,
+			expectedAttrs: []attribute.KeyValue{
+				attribute.String("gen_ai.image.count", "1"),
+				attribute.String("gen_ai.image.urls", "https://example.com/img.png"),
+			},
+			expectedStatus: trace.Status{Code: codes.Ok, Description: ""},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			recorder := NewImageGenerationRecorder(nil)
+
+			resp := &openaisdk.ImagesResponse{}
+			err := json.Unmarshal(tt.respBody, resp)
+			require.NoError(t, err)
+
+			actualSpan := testotel.RecordWithSpan(t, func(span oteltrace.Span) bool {
+				recorder.RecordResponse(span, resp)
+				return false
+			})
+
+			// Check that key attributes are present
+			attrs := attributesToMap(actualSpan.Attributes)
+			require.Equal(t, "1", attrs["gen_ai.image.count"])
+			require.Equal(t, "https://example.com/img.png", attrs["gen_ai.image.urls"])
+			require.Equal(t, trace.Status{Code: codes.Ok, Description: ""}, actualSpan.Status)
+		})
 	}
-	recorder.RecordResponse(span, resp)
-	span.End()
-
-	spans := exporter.GetSpans()
-	require.Len(t, spans, 1)
-	got := spans[0]
-	attrs := attributesToMap(got.Attributes)
-	// Count and urls should be present
-	require.Equal(t, "1", attrs["gen_ai.image.count"])
-	require.Equal(t, "https://example.com/img.png", attrs["gen_ai.image.urls"])
+}
+
+func TestImageGenerationRecorder_RecordResponseOnError(t *testing.T) {
+	recorder := NewImageGenerationRecorder(nil)
+
+	actualSpan := testotel.RecordWithSpan(t, func(span oteltrace.Span) bool {
+		recorder.RecordResponseOnError(span, 400, []byte(`{"error":{"message":"Invalid request","type":"invalid_request_error"}}`))
+		return false
+	})
+
+	require.Equal(t, trace.Status{
+		Code:        codes.Error,
+		Description: `Error code: 400 - {"error":{"message":"Invalid request","type":"invalid_request_error"}}`,
+	}, actualSpan.Status)
 }
 
 // attributesToMap converts attribute KeyValue to a simple map for assertions.