From 9fba63b0b31ba0d30eabce395d06249c82bbeec2 Mon Sep 17 00:00:00 2001 From: noalimoy Date: Tue, 10 Mar 2026 16:07:36 +0200 Subject: [PATCH] test(bbr): add request plugin header mutation tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add unit and integration tests to verify SetHeader/RemoveHeader interactions across multiple request plugins in all combinations. Changes: - Add request_plugins_test.go with 8 unit test scenarios: set→remove, remove→set, double-remove (idempotent), set same value (no-op optimization), last-writer-wins, and mixed header operations - Add plugins_hermetic_test.go with matching 8 integration test scenarios over a real gRPC ext_proc server - Extend NewBBRHarness to accept custom request plugins (nil = default BodyFieldToHeaderPlugin) - Update bbr_success_total metric assertion to account for shared global Prometheus counter across test files Signed-off-by: noalimoy --- pkg/bbr/handlers/request_test.go | 299 ++++++++++++++++++++++++++++++- 1 file changed, 298 insertions(+), 1 deletion(-) diff --git a/pkg/bbr/handlers/request_test.go b/pkg/bbr/handlers/request_test.go index b40ac70576..22905cc51b 100644 --- a/pkg/bbr/handlers/request_test.go +++ b/pkg/bbr/handlers/request_test.go @@ -41,6 +41,8 @@ import ( const modelField = "model" +// === Request Headers Tests === + func TestHandleRequestHeaders(t *testing.T) { tests := []struct { name string @@ -196,6 +198,8 @@ func TestHandleRequestHeaders(t *testing.T) { } } +// === Request Body Tests (built-in plugins) === + func TestHandleRequestBody(t *testing.T) { metrics.Register() ctx := logutil.NewTestLoggerIntoContext(context.Background()) @@ -529,7 +533,7 @@ func TestHandleRequestBody(t *testing.T) { }) } - // Assert BBR metrics: 2 model not in body, 1 model empty string, 4 successful model-from-body cases. + // Assert BBR metrics: 2 model not in body, 1 model empty string, 7 successful model-from-body cases. wantMetrics := ` # HELP bbr_body_field_empty_total [ALPHA] Count of times a field was found in a request body but was empty. # TYPE bbr_body_field_empty_total counter @@ -594,6 +598,299 @@ func TestHandleRequestBodyWithPluginMetrics(t *testing.T) { } } +// === Request Body Tests (multi-plugin header mutations) === + +// fakeRequestPlugin implements framework.RequestProcessor for testing +// multi-plugin header mutation scenarios. +type fakeRequestPlugin struct { + name string + mutateFn func(ctx context.Context, request *framework.InferenceRequest) error +} + +func (p *fakeRequestPlugin) TypedName() epp.TypedName { + return epp.TypedName{Type: "fake", Name: p.name} +} + +func (p *fakeRequestPlugin) ProcessRequest(ctx context.Context, _ *framework.CycleState, request *framework.InferenceRequest) error { + return p.mutateFn(ctx, request) +} + +var _ framework.RequestProcessor = &fakeRequestPlugin{} + +// TestHandleRequestBody_MultiPluginHeaderMutations tests the end-to-end behavior of +// HandleRequestBody when multiple request plugins set and/or remove headers. +// Each sub-test verifies the HeaderMutation in the resulting ProcessingResponse. +func TestHandleRequestBody_MultiPluginHeaderMutations(t *testing.T) { + ctx := logutil.NewTestLoggerIntoContext(context.Background()) + + tests := []struct { + name string + plugins []framework.RequestProcessor + initialHeaders map[string]string + wantSetHeaders map[string]string + wantRemoved []string + }{ + { + // Plugin1 adds X-Custom, Plugin2 removes it. + // The header was never in the original Envoy request, so Envoy treats + // the removal as a no-op. The net visible effect is: nothing changed. + // However, RemoveHeader() does record it in removedHeaders because + // the key existed in Headers at the time of removal. + name: "set then remove same header - cancels out", + plugins: []framework.RequestProcessor{ + &fakeRequestPlugin{ + name: "setter", + mutateFn: func(_ context.Context, req *framework.InferenceRequest) error { + req.SetHeader("X-Custom", "value1") + return nil + }, + }, + &fakeRequestPlugin{ + name: "remover", + mutateFn: func(_ context.Context, req *framework.InferenceRequest) error { + req.RemoveHeader("X-Custom") + return nil + }, + }, + }, + wantSetHeaders: map[string]string{}, + wantRemoved: []string{"X-Custom"}, + }, + { + // Plugin1 adds a new header, Plugin2 removes a pre-existing one. + // Both mutations should appear in the response. + name: "set then remove different headers - both apply", + plugins: []framework.RequestProcessor{ + &fakeRequestPlugin{ + name: "setter", + mutateFn: func(_ context.Context, req *framework.InferenceRequest) error { + req.SetHeader("X-New", "hello") + return nil + }, + }, + &fakeRequestPlugin{ + name: "remover", + mutateFn: func(_ context.Context, req *framework.InferenceRequest) error { + req.RemoveHeader("X-Existing") + return nil + }, + }, + }, + initialHeaders: map[string]string{ + "X-Existing": "old-value", + }, + wantSetHeaders: map[string]string{ + "X-New": "hello", + }, + wantRemoved: []string{"X-Existing"}, + }, + { + // RemoveHeader on a key that was never in Headers is a no-op: + // the guard `if _, ok := r.Headers[key]; ok` prevents any mutation. + name: "remove non-existing header - no-op", + plugins: []framework.RequestProcessor{ + &fakeRequestPlugin{ + name: "remover", + mutateFn: func(_ context.Context, req *framework.InferenceRequest) error { + req.RemoveHeader("X-Ghost") + return nil + }, + }, + }, + wantSetHeaders: map[string]string{}, + wantRemoved: nil, + }, + { + // Plugin1 removes a pre-existing header, Plugin2 re-sets it. + // SetHeader clears the key from removedHeaders, so the final result + // is a set with the new value and no removal. + name: "remove then set same header - new value wins", + plugins: []framework.RequestProcessor{ + &fakeRequestPlugin{ + name: "remover", + mutateFn: func(_ context.Context, req *framework.InferenceRequest) error { + req.RemoveHeader("X-Reuse") + return nil + }, + }, + &fakeRequestPlugin{ + name: "setter", + mutateFn: func(_ context.Context, req *framework.InferenceRequest) error { + req.SetHeader("X-Reuse", "new-value") + return nil + }, + }, + }, + initialHeaders: map[string]string{ + "X-Reuse": "old-value", + }, + wantSetHeaders: map[string]string{ + "X-Reuse": "new-value", + }, + wantRemoved: nil, + }, + { + // Both plugins set the same header key. Plugins run sequentially, + // so the last writer wins. + name: "two plugins set same header - last wins", + plugins: []framework.RequestProcessor{ + &fakeRequestPlugin{ + name: "setter1", + mutateFn: func(_ context.Context, req *framework.InferenceRequest) error { + req.SetHeader("X-Shared", "first") + return nil + }, + }, + &fakeRequestPlugin{ + name: "setter2", + mutateFn: func(_ context.Context, req *framework.InferenceRequest) error { + req.SetHeader("X-Shared", "second") + return nil + }, + }, + }, + wantSetHeaders: map[string]string{ + "X-Shared": "second", + }, + wantRemoved: nil, + }, + { + // Two plugins set different header keys. Both should appear in the response. + name: "two plugins set different headers - both apply", + plugins: []framework.RequestProcessor{ + &fakeRequestPlugin{ + name: "setter-a", + mutateFn: func(_ context.Context, req *framework.InferenceRequest) error { + req.SetHeader("X-First", "aaa") + return nil + }, + }, + &fakeRequestPlugin{ + name: "setter-b", + mutateFn: func(_ context.Context, req *framework.InferenceRequest) error { + req.SetHeader("X-Second", "bbb") + return nil + }, + }, + }, + wantSetHeaders: map[string]string{ + "X-First": "aaa", + "X-Second": "bbb", + }, + wantRemoved: nil, + }, + { + // Two plugins both remove the same pre-existing header. + // The second RemoveHeader is a no-op because the header is already gone. + // The header should appear exactly once in removedHeaders. + name: "two plugins remove same header - idempotent", + plugins: []framework.RequestProcessor{ + &fakeRequestPlugin{ + name: "remover1", + mutateFn: func(_ context.Context, req *framework.InferenceRequest) error { + req.RemoveHeader("X-Dup") + return nil + }, + }, + &fakeRequestPlugin{ + name: "remover2", + mutateFn: func(_ context.Context, req *framework.InferenceRequest) error { + req.RemoveHeader("X-Dup") + return nil + }, + }, + }, + initialHeaders: map[string]string{ + "X-Dup": "value", + }, + wantSetHeaders: map[string]string{}, + wantRemoved: []string{"X-Dup"}, + }, + { + // A plugin sets a header to the same value it already has. + // The SetHeader optimization (types.go:43) should skip the mutation. + name: "set existing header to same value - no mutation", + plugins: []framework.RequestProcessor{ + &fakeRequestPlugin{ + name: "noop-setter", + mutateFn: func(_ context.Context, req *framework.InferenceRequest) error { + req.SetHeader("X-Keep", "original") + return nil + }, + }, + }, + initialHeaders: map[string]string{ + "X-Keep": "original", + }, + wantSetHeaders: map[string]string{}, + wantRemoved: nil, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + server := NewServer(false, tc.plugins, []framework.ResponseProcessor{}) + reqCtx := &RequestContext{ + Request: framework.NewInferenceRequest(), + CycleState: framework.NewCycleState(), + } + for k, v := range tc.initialHeaders { + reqCtx.Request.Headers[k] = v + } + + bodyBytes, err := json.Marshal(map[string]any{"model": "test-model", "prompt": "test"}) + if err != nil { + t.Fatalf("Failed to marshal request body: %v", err) + } + + resp, err := server.HandleRequestBody(ctx, reqCtx, bodyBytes) + if err != nil { + t.Fatalf("HandleRequestBody returned unexpected error: %v", err) + } + + want := buildNonStreamingResponse(tc.wantSetHeaders, tc.wantRemoved) + envoytest.SortSetHeadersInResponses(want) + envoytest.SortSetHeadersInResponses(resp) + + if diff := cmp.Diff(want, resp, protocmp.Transform()); diff != "" { + t.Errorf("HandleRequestBody response mismatch (-want +got):\n%s", diff) + } + }) + } +} + +// buildNonStreamingResponse constructs the expected ProcessingResponse for a +// non-streaming HandleRequestBody call with the given header mutations. +func buildNonStreamingResponse(setHeaders map[string]string, removeHeaders []string) []*extProcPb.ProcessingResponse { + setHeaderOpts := make([]*basepb.HeaderValueOption, 0, len(setHeaders)) + for k, v := range setHeaders { + setHeaderOpts = append(setHeaderOpts, &basepb.HeaderValueOption{ + Header: &basepb.HeaderValue{ + Key: k, + RawValue: []byte(v), + }, + }) + } + + return []*extProcPb.ProcessingResponse{ + { + Response: &extProcPb.ProcessingResponse_RequestBody{ + RequestBody: &extProcPb.BodyResponse{ + Response: &extProcPb.CommonResponse{ + ClearRouteCache: true, + HeaderMutation: &extProcPb.HeaderMutation{ + SetHeaders: setHeaderOpts, + RemoveHeaders: removeHeaders, + }, + }, + }, + }, + }, + } +} + +// === Request Body Tests (body mutations) === + type bodyMutatingPlugin struct { name string mutateFn func(ctx context.Context, cycleState *framework.CycleState, request *framework.InferenceRequest) error