Skip to content

Commit 87a9506

Browse files
mathetakemathetake
committed
test build
Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>
1 parent 32ff20e commit 87a9506

File tree

8 files changed

+101
-120
lines changed

8 files changed

+101
-120
lines changed

‎internal/extproc/chatcompletion_processor_test.go‎

Lines changed: 19 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -256,13 +256,11 @@ func Test_chatCompletionProcessorUpstreamFilter_ProcessResponseBody(t *testing.T
256256
})
257257
t.Run("ok", func(t *testing.T) {
258258
inBody := &extprocv3.HttpBody{Body: []byte("some-body"), EndOfStream: true}
259-
expBodyMut := &extprocv3.BodyMutation{}
260-
expHeadMut := &extprocv3.HeaderMutation{}
261259
mm := &mockChatCompletionMetrics{}
262260
mt := &mockTranslator{
263261
t: t, expResponseBody: inBody,
264-
retBodyMutation: expBodyMut, retHeaderMutation: expHeadMut,
265-
retUsedToken: translator.LLMTokenUsage{OutputTokens: 123, InputTokens: 1, CachedInputTokens: 1},
262+
retHeaderMutation: []internalapi.Header{{"foo", "bar"}},
263+
retUsedToken: translator.LLMTokenUsage{OutputTokens: 123, InputTokens: 1, CachedInputTokens: 1},
266264
}
267265

268266
celProgInt, err := llmcostcel.NewProgram("54321")
@@ -297,8 +295,10 @@ func Test_chatCompletionProcessorUpstreamFilter_ProcessResponseBody(t *testing.T
297295
res, err := p.ProcessResponseBody(t.Context(), inBody)
298296
require.NoError(t, err)
299297
commonRes := res.Response.(*extprocv3.ProcessingResponse_ResponseBody).ResponseBody.Response
300-
require.Equal(t, expBodyMut, commonRes.BodyMutation)
301-
require.Equal(t, expHeadMut, commonRes.HeaderMutation)
298+
require.Nil(t, commonRes.BodyMutation)
299+
require.Len(t, commonRes.HeaderMutation.SetHeaders, 1)
300+
require.Equal(t, "foo", commonRes.HeaderMutation.SetHeaders[0].Header.Key)
301+
require.Equal(t, "bar", commonRes.HeaderMutation.SetHeaders[0].Header.Value)
302302
mm.RequireRequestSuccess(t)
303303
require.Equal(t, 124, mm.tokenUsageCount) // 1 input + 123 output
304304

@@ -321,8 +321,8 @@ func Test_chatCompletionProcessorUpstreamFilter_ProcessResponseBody(t *testing.T
321321
// Verify we record failure for non-2xx responses and do it exactly once (defer suppressed).
322322
t.Run("non-2xx status failure once", func(t *testing.T) {
323323
inBody := &extprocv3.HttpBody{Body: []byte("error-body"), EndOfStream: true}
324-
expHeadMut := &extprocv3.HeaderMutation{}
325-
expBodyMut := &extprocv3.BodyMutation{}
324+
expHeadMut := []internalapi.Header{{"foo", "bar"}}
325+
expBodyMut := []byte("error-body")
326326
mm := &mockChatCompletionMetrics{}
327327
mt := &mockTranslator{t: t, expResponseBody: inBody, retHeaderMutation: expHeadMut, retBodyMutation: expBodyMut}
328328
p := &chatCompletionProcessorUpstreamFilter{
@@ -333,7 +333,10 @@ func Test_chatCompletionProcessorUpstreamFilter_ProcessResponseBody(t *testing.T
333333
res, err := p.ProcessResponseBody(t.Context(), inBody)
334334
require.NoError(t, err)
335335
commonRes := res.Response.(*extprocv3.ProcessingResponse_ResponseBody).ResponseBody.Response
336-
require.Equal(t, expBodyMut, commonRes.BodyMutation)
336+
require.Equal(t, "error-body", commonRes.BodyMutation.GetBody())
337+
require.Len(t, commonRes.HeaderMutation.SetHeaders, 1)
338+
require.Equal(t, "foo", commonRes.HeaderMutation.SetHeaders[0].Header.Key)
339+
require.Equal(t, []byte("bar"), commonRes.HeaderMutation.SetHeaders[0].Header.RawValue)
337340
require.Equal(t, expHeadMut, commonRes.HeaderMutation)
338341
mm.RequireRequestFailure(t)
339342
})
@@ -469,16 +472,15 @@ func Test_chatCompletionProcessorUpstreamFilter_ProcessRequestHeaders(t *testing
469472
t.Run("ok", func(t *testing.T) {
470473
someBody := bodyFromModel(t, "some-model", tc.stream, nil)
471474
headers := map[string]string{":path": "/foo", internalapi.ModelNameHeaderKeyDefault: "some-model"}
472-
headerMut := &extprocv3.HeaderMutation{
473-
SetHeaders: []*corev3.HeaderValueOption{{Header: &corev3.HeaderValue{Key: "foo", RawValue: []byte("bar")}}},
474-
}
475-
bodyMut := &extprocv3.BodyMutation{Mutation: &extprocv3.BodyMutation_Body{Body: []byte("some body")}}
475+
headerMut := []internalapi.Header{{"foo", "bar"}}
476+
bodyMut := []byte("some body")
476477

477478
var expBody openai.ChatCompletionRequest
478479
require.NoError(t, json.Unmarshal(someBody, &expBody))
479480
if tc.stream && tc.forcedIncludeUsage {
480481
expBody.StreamOptions = &openai.StreamOptions{IncludeUsage: true}
481482
}
483+
482484
mt := mockTranslator{
483485
t: t, expRequestBody: &expBody, retHeaderMutation: headerMut,
484486
retBodyMutation: bodyMut, expForceRequestBodyMutation: tc.forcedIncludeUsage,
@@ -501,8 +503,10 @@ func Test_chatCompletionProcessorUpstreamFilter_ProcessRequestHeaders(t *testing
501503
require.Equal(t, mt, p.translator)
502504
require.NotNil(t, resp)
503505
commonRes := resp.Response.(*extprocv3.ProcessingResponse_RequestHeaders).RequestHeaders.Response
504-
require.Equal(t, headerMut, commonRes.HeaderMutation)
505-
require.Equal(t, bodyMut, commonRes.BodyMutation)
506+
require.Equal(t, string(bodyMut), string(commonRes.BodyMutation.GetBody()))
507+
require.Len(t, commonRes.HeaderMutation.SetHeaders, 1)
508+
require.Equal(t, "foo", commonRes.HeaderMutation.SetHeaders[0].Header.Key)
509+
require.Equal(t, []byte("bar"), commonRes.HeaderMutation.SetHeaders[0].Header.RawValue)
506510

507511
mm.RequireRequestNotCompleted(t)
508512
// Verify models were set

‎internal/extproc/completions_processor_test.go‎

Lines changed: 23 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -131,11 +131,7 @@ func Test_completionsProcessorUpstreamFilter_ProcessResponseHeaders(t *testing.T
131131
const headerName = ":test-header:"
132132
const headerValue = ":test-header-value:"
133133
mt.expHeaders[headerName] = headerValue
134-
mt.resHeaderMutation = &extprocv3.HeaderMutation{
135-
SetHeaders: []*corev3.HeaderValueOption{
136-
{Header: &corev3.HeaderValue{Key: headerName, RawValue: []byte(headerValue)}},
137-
},
138-
}
134+
mt.resHeaderMutation = []internalapi.Header{{headerName, headerValue}}
139135

140136
resp, err := p.ProcessResponseHeaders(t.Context(), &corev3.HeaderMap{
141137
Headers: []*corev3.HeaderValue{
@@ -165,14 +161,8 @@ func Test_completionsProcessorUpstreamFilter_ProcessResponseBody(t *testing.T) {
165161
metrics: mm,
166162
}
167163

168-
mt.resErrorHeaderMutation = &extprocv3.HeaderMutation{
169-
SetHeaders: []*corev3.HeaderValueOption{
170-
{Header: &corev3.HeaderValue{Key: "test", RawValue: []byte("error")}},
171-
},
172-
}
173-
mt.resErrorBodyMutation = &extprocv3.BodyMutation{
174-
Mutation: &extprocv3.BodyMutation_Body{Body: []byte("error body")},
175-
}
164+
mt.resErrorHeaderMutation = []internalapi.Header{{"test", "error"}}
165+
mt.resErrorBodyMutation = []byte("test error")
176166

177167
resp, err := p.ProcessResponseBody(t.Context(), &extprocv3.HttpBody{Body: []byte("test error")})
178168
require.NoError(t, err)
@@ -197,14 +187,8 @@ func Test_completionsProcessorUpstreamFilter_ProcessResponseBody(t *testing.T) {
197187
metrics: mm,
198188
}
199189

200-
mt.resHeaderMutation = &extprocv3.HeaderMutation{
201-
SetHeaders: []*corev3.HeaderValueOption{
202-
{Header: &corev3.HeaderValue{Key: "test", RawValue: []byte("success")}},
203-
},
204-
}
205-
mt.resBodyMutation = &extprocv3.BodyMutation{
206-
Mutation: &extprocv3.BodyMutation_Body{Body: []byte("response body")},
207-
}
190+
mt.resHeaderMutation = []internalapi.Header{{"test", "success"}}
191+
mt.resBodyMutation = []byte("response body")
208192
mt.resTokenUsage = translator.LLMTokenUsage{
209193
InputTokens: 10,
210194
OutputTokens: 20,
@@ -404,10 +388,10 @@ func completionBodyFromModel(t *testing.T, model string) []byte {
404388
type mockCompletionTranslator struct {
405389
t *testing.T
406390
expHeaders map[string]string
407-
resHeaderMutation *extprocv3.HeaderMutation
408-
resBodyMutation *extprocv3.BodyMutation
409-
resErrorHeaderMutation *extprocv3.HeaderMutation
410-
resErrorBodyMutation *extprocv3.BodyMutation
391+
resHeaderMutation []internalapi.Header
392+
resBodyMutation []byte
393+
resErrorHeaderMutation []internalapi.Header
394+
resErrorBodyMutation []byte
411395
resTokenUsage translator.LLMTokenUsage
412396
resModel internalapi.ResponseModel
413397
err error
@@ -649,10 +633,10 @@ func Test_completionsProcessorUpstreamFilter_ProcessResponseBody_NonSuccess(t *t
649633
// Verify we record failure for non-2xx responses and do it exactly once (defer suppressed).
650634
t.Run("non-2xx status failure once", func(t *testing.T) {
651635
inBody := &extprocv3.HttpBody{Body: []byte("error-body"), EndOfStream: true}
652-
expHeadMut := &extprocv3.HeaderMutation{}
653-
expBodyMut := &extprocv3.BodyMutation{}
654636
mm := &mockCompletionMetrics{}
655-
mt := &mockCompletionTranslator{t: t, resErrorHeaderMutation: expHeadMut, resErrorBodyMutation: expBodyMut}
637+
mt := &mockCompletionTranslator{t: t, resErrorHeaderMutation: []internalapi.Header{
638+
{"foo", "bar"},
639+
}, resErrorBodyMutation: []byte("abcd")}
656640
p := &completionsProcessorUpstreamFilter{
657641
translator: mt,
658642
metrics: mm,
@@ -661,8 +645,10 @@ func Test_completionsProcessorUpstreamFilter_ProcessResponseBody_NonSuccess(t *t
661645
res, err := p.ProcessResponseBody(t.Context(), inBody)
662646
require.NoError(t, err)
663647
commonRes := res.Response.(*extprocv3.ProcessingResponse_ResponseBody).ResponseBody.Response
664-
require.Equal(t, expBodyMut, commonRes.BodyMutation)
665-
require.Equal(t, expHeadMut, commonRes.HeaderMutation)
648+
require.Equal(t, "abcd", string(commonRes.BodyMutation.GetBody()))
649+
require.Len(t, commonRes.HeaderMutation.SetHeaders, 1)
650+
require.Equal(t, "foo", commonRes.HeaderMutation.SetHeaders[0].Header.Key)
651+
require.Equal(t, "bar", string(commonRes.HeaderMutation.SetHeaders[0].Header.RawValue))
666652
mm.RequireRequestFailure(t)
667653
})
668654
}
@@ -767,13 +753,12 @@ func Test_completionsProcessorUpstreamFilter_CELCostEvaluation(t *testing.T) {
767753
// Using exactly the same test data as chat completion CEL test
768754
t.Run("CEL expressions with token usage", func(t *testing.T) {
769755
inBody := &extprocv3.HttpBody{Body: []byte("response-body"), EndOfStream: true}
770-
expBodyMut := &extprocv3.BodyMutation{}
771-
expHeadMut := &extprocv3.HeaderMutation{}
756+
expBody := []byte("response-body")
772757
mm := &mockCompletionMetrics{}
773758
mt := &mockCompletionTranslator{
774759
t: t,
775-
resBodyMutation: expBodyMut,
776-
resHeaderMutation: expHeadMut,
760+
resBodyMutation: expBody,
761+
resHeaderMutation: []internalapi.Header{{"foo", "bar"}},
777762
resTokenUsage: translator.LLMTokenUsage{
778763
OutputTokens: 123,
779764
InputTokens: 1,
@@ -813,8 +798,10 @@ func Test_completionsProcessorUpstreamFilter_CELCostEvaluation(t *testing.T) {
813798
res, err := p.ProcessResponseBody(t.Context(), inBody)
814799
require.NoError(t, err)
815800
commonRes := res.Response.(*extprocv3.ProcessingResponse_ResponseBody).ResponseBody.Response
816-
require.Equal(t, expBodyMut, commonRes.BodyMutation)
817-
require.Equal(t, expHeadMut, commonRes.HeaderMutation)
801+
require.Equal(t, string(expBody), string(commonRes.BodyMutation.GetBody()))
802+
require.Len(t, commonRes.HeaderMutation, 1)
803+
require.Equal(t, "foo", commonRes.HeaderMutation.SetHeaders[0].Header.Key)
804+
require.Equal(t, "bar", string(commonRes.HeaderMutation.SetHeaders[0].Header.RawValue))
818805
mm.RequireRequestSuccess(t)
819806
require.Equal(t, 124, mm.tokenUsageCount) // 1 input + 123 output
820807
md := res.DynamicMetadata

‎internal/extproc/embeddings_processor_test.go‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -145,8 +145,8 @@ func Test_embeddingsProcessorUpstreamFilter_ProcessResponseBody(t *testing.T) {
145145
})
146146
t.Run("ok", func(t *testing.T) {
147147
inBody := &extprocv3.HttpBody{Body: []byte("some-body"), EndOfStream: true}
148-
expBodyMut := &extprocv3.BodyMutation{}
149-
expHeadMut := &extprocv3.HeaderMutation{}
148+
expBodyMut := []byte("some body")
149+
expHeadMut := []internalapi.Header{{"foo", "bar"}}
150150
mm := &mockEmbeddingsMetrics{}
151151
mt := &mockEmbeddingTranslator{
152152
t: t, expResponseBody: inBody,
@@ -184,8 +184,10 @@ func Test_embeddingsProcessorUpstreamFilter_ProcessResponseBody(t *testing.T) {
184184
res, err := p.ProcessResponseBody(t.Context(), inBody)
185185
require.NoError(t, err)
186186
commonRes := res.Response.(*extprocv3.ProcessingResponse_ResponseBody).ResponseBody.Response
187-
require.Equal(t, expBodyMut, commonRes.BodyMutation)
188-
require.Equal(t, expHeadMut, commonRes.HeaderMutation)
187+
require.Equal(t, string(expBodyMut), string(commonRes.BodyMutation.GetBody()))
188+
require.Len(t, commonRes.HeaderMutation.SetHeaders, 1)
189+
require.Equal(t, "foo", commonRes.HeaderMutation.SetHeaders[0].Header.Key)
190+
require.Equal(t, []byte("bar"), commonRes.HeaderMutation.SetHeaders[0].Header.RawValue)
189191
mm.RequireRequestSuccess(t)
190192
mm.RequireTokenUsage(t, 123)
191193

@@ -327,10 +329,8 @@ func Test_embeddingsProcessorUpstreamFilter_ProcessRequestHeaders(t *testing.T)
327329
t.Run("ok", func(t *testing.T) {
328330
someBody := embeddingBodyFromModel(t, "some-model")
329331
headers := map[string]string{":path": "/foo", internalapi.ModelNameHeaderKeyDefault: "some-model"}
330-
headerMut := &extprocv3.HeaderMutation{
331-
SetHeaders: []*corev3.HeaderValueOption{{Header: &corev3.HeaderValue{Key: "foo", RawValue: []byte("bar")}}},
332-
}
333-
bodyMut := &extprocv3.BodyMutation{Mutation: &extprocv3.BodyMutation_Body{Body: []byte("some body")}}
332+
headerMut := []internalapi.Header{{"foo", "bar"}}
333+
bodyMut := []byte("some body")
334334

335335
var expBody openai.EmbeddingRequest
336336
require.NoError(t, json.Unmarshal(someBody, &expBody))
@@ -351,8 +351,10 @@ func Test_embeddingsProcessorUpstreamFilter_ProcessRequestHeaders(t *testing.T)
351351
require.Equal(t, mt, p.translator)
352352
require.NotNil(t, resp)
353353
commonRes := resp.Response.(*extprocv3.ProcessingResponse_RequestHeaders).RequestHeaders.Response
354-
require.Equal(t, headerMut, commonRes.HeaderMutation)
355-
require.Equal(t, bodyMut, commonRes.BodyMutation)
354+
require.Len(t, commonRes.HeaderMutation.SetHeaders, 1)
355+
require.Equal(t, "foo", commonRes.HeaderMutation.SetHeaders[0].Header.Key)
356+
require.Equal(t, []byte("bar"), commonRes.HeaderMutation.SetHeaders[0].Header.RawValue)
357+
require.Equal(t, string(bodyMut), string(commonRes.BodyMutation.GetBody()))
356358

357359
mm.RequireRequestNotCompleted(t)
358360
// Verify models were set

‎internal/extproc/imagegeneration_processor_test.go‎

Lines changed: 18 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -238,12 +238,9 @@ func Test_imageGenerationProcessorUpstreamFilter_ProcessResponseBody(t *testing.
238238
})
239239
t.Run("ok", func(t *testing.T) {
240240
inBody := &extprocv3.HttpBody{Body: []byte("some-body"), EndOfStream: true}
241-
expBodyMut := &extprocv3.BodyMutation{}
242-
expHeadMut := &extprocv3.HeaderMutation{}
243241
mm := &mockImageGenerationMetrics{}
244242
mt := &mockImageGenerationTranslator{
245243
t: t, expResponseBody: inBody,
246-
retBodyMutation: expBodyMut, retHeaderMutation: expHeadMut,
247244
retUsedToken: translator.LLMTokenUsage{OutputTokens: 123, InputTokens: 1},
248245
}
249246

@@ -277,8 +274,8 @@ func Test_imageGenerationProcessorUpstreamFilter_ProcessResponseBody(t *testing.
277274
res, err := p.ProcessResponseBody(t.Context(), inBody)
278275
require.NoError(t, err)
279276
commonRes := res.Response.(*extprocv3.ProcessingResponse_ResponseBody).ResponseBody.Response
280-
require.Equal(t, expBodyMut, commonRes.BodyMutation)
281-
require.Equal(t, expHeadMut, commonRes.HeaderMutation)
277+
require.Equal(t, nil, commonRes.BodyMutation)
278+
require.Equal(t, &extprocv3.HeaderMutation{}, commonRes.HeaderMutation)
282279
mm.RequireRequestSuccess(t)
283280
require.Equal(t, 124, mm.tokenUsageCount) // 1 input + 123 output
284281

@@ -299,10 +296,11 @@ func Test_imageGenerationProcessorUpstreamFilter_ProcessResponseBody(t *testing.
299296
// Verify we record failure for non-2xx responses and do it exactly once (defer suppressed), and span records error.
300297
t.Run("non-2xx status failure once", func(t *testing.T) {
301298
inBody := &extprocv3.HttpBody{Body: []byte("error-body"), EndOfStream: true}
302-
expHeadMut := &extprocv3.HeaderMutation{}
303-
expBodyMut := &extprocv3.BodyMutation{}
299+
expBody := []byte("error-body")
304300
mm := &mockImageGenerationMetrics{}
305-
mt := &mockImageGenerationTranslator{t: t, expResponseBody: inBody, retHeaderMutation: expHeadMut, retBodyMutation: expBodyMut}
301+
mt := &mockImageGenerationTranslator{t: t, expResponseBody: inBody, retHeaderMutation: []internalapi.Header{
302+
{"foo", "bar"},
303+
}, retBodyMutation: expBody}
306304
p := &imageGenerationProcessorUpstreamFilter{
307305
translator: mt,
308306
metrics: mm,
@@ -313,8 +311,10 @@ func Test_imageGenerationProcessorUpstreamFilter_ProcessResponseBody(t *testing.
313311
res, err := p.ProcessResponseBody(t.Context(), inBody)
314312
require.NoError(t, err)
315313
commonRes := res.Response.(*extprocv3.ProcessingResponse_ResponseBody).ResponseBody.Response
316-
require.Equal(t, expBodyMut, commonRes.BodyMutation)
317-
require.Equal(t, expHeadMut, commonRes.HeaderMutation)
314+
require.Equal(t, expBody, commonRes.BodyMutation.GetBody())
315+
require.Len(t, commonRes.HeaderMutation.SetHeaders, 1)
316+
require.Equal(t, "foo", commonRes.HeaderMutation.SetHeaders[0].Header.Key)
317+
require.Equal(t, "bar", commonRes.HeaderMutation.SetHeaders[0].Header.Value)
318318
mm.RequireRequestFailure(t)
319319
// assert span error recorded
320320
s := p.span.(*mockImageGenerationSpan)
@@ -333,8 +333,8 @@ func Test_imageGenerationProcessorUpstreamFilter_ProcessResponseBody(t *testing.
333333
mm := &mockImageGenerationMetrics{}
334334
mt := &mockImageGenerationTranslator{
335335
// translator returns a non-nil body mutation indicating processor changed body
336-
retBodyMutation: &extprocv3.BodyMutation{Mutation: &extprocv3.BodyMutation_Body{Body: []byte("changed")}},
337-
retHeaderMutation: &extprocv3.HeaderMutation{},
336+
retBodyMutation: []byte("changed"),
337+
retHeaderMutation: []internalapi.Header{{"foo", "bar"}},
338338
}
339339
p := &imageGenerationProcessorUpstreamFilter{
340340
translator: mt,
@@ -349,6 +349,10 @@ func Test_imageGenerationProcessorUpstreamFilter_ProcessResponseBody(t *testing.
349349
commonRes := res.Response.(*extprocv3.ProcessingResponse_ResponseBody).ResponseBody.Response
350350
reqHM := commonRes.HeaderMutation
351351
require.Contains(t, reqHM.RemoveHeaders, "content-encoding")
352+
require.Len(t, reqHM.SetHeaders, 1)
353+
require.Equal(t, "foo", reqHM.SetHeaders[0].Header.Key)
354+
require.Equal(t, "bar", reqHM.SetHeaders[0].Header.Value)
355+
require.Equal(t, []byte("changed"), commonRes.BodyMutation.GetBody())
352356
mm.RequireRequestSuccess(t)
353357
})
354358
}
@@ -470,8 +474,8 @@ type mockImageGenerationTranslator struct {
470474
expHeaders map[string]string
471475
expForceRequestBodyMutation bool
472476
retErr error
473-
retHeaderMutation *extprocv3.HeaderMutation
474-
retBodyMutation *extprocv3.BodyMutation
477+
retHeaderMutation []internalapi.Header
478+
retBodyMutation []byte
475479
retUsedToken translator.LLMTokenUsage
476480
retResponseModel internalapi.ResponseModel
477481
}

0 commit comments

Comments
 (0)