Add integration tests for response plugin execution

Follow-up to kubernetes-sigs#2354 / PR kubernetes-sigs#2369. Exercises the full Process loop
(request body -> response headers -> response body) to validate that
response plugins can mutate the response body over a real gRPC
ext_proc stream.

Covers:
- Unary mode: response plugin mutates the response body
- Unary mode: no response plugins results in passthrough behavior

Signed-off-by: Abdallah Samara <[email protected]>
Made-with: Cursor

Author: abdallahsamabd

test/integration/bbr/body_mutation_test.go

@@ -60,7 +60,7 @@ func TestBodyMutation_Unary(t *testing.T) {
 
 	plugin := &bodyMutatingPlugin{fieldName: "injected", fieldValue: "test-value"}
 	baseModelToHeaderPlugin := &basemodelextractor.BaseModelToHeaderPlugin{AdaptersStore: basemodelextractor.NewAdaptersStore()}
-	h := NewBBRHarnessWithPlugins(t, ctx, false, []framework.RequestProcessor{plugin, baseModelToHeaderPlugin})
+	h := NewBBRHarnessWithPlugins(t, ctx, false, []framework.RequestProcessor{plugin, baseModelToHeaderPlugin}, []framework.ResponseProcessor{})
 
 	body := map[string]any{"prompt": "hello"}
 	bodyBytes, _ := json.Marshal(body)
@@ -127,7 +127,7 @@ func TestBodyMutation_Streaming(t *testing.T) {
 
 	plugin := &bodyMutatingPlugin{fieldName: "injected", fieldValue: "test-value"}
 	baseModelToHeaderPlugin := &basemodelextractor.BaseModelToHeaderPlugin{AdaptersStore: basemodelextractor.NewAdaptersStore()}
-	h := NewBBRHarnessWithPlugins(t, ctx, true, []framework.RequestProcessor{plugin, baseModelToHeaderPlugin})
+	h := NewBBRHarnessWithPlugins(t, ctx, true, []framework.RequestProcessor{plugin, baseModelToHeaderPlugin}, []framework.ResponseProcessor{})
 
 	body := map[string]any{"prompt": "hello"}
 	bodyBytes, _ := json.Marshal(body)

test/integration/bbr/harness.go

@@ -52,7 +52,7 @@ type BBRHarness struct {
 }
 
 // NewBBRHarness boots up an isolated BBR server on a random port with the default
-// BodyFieldToHeaderPlugin for model extraction.
+// BodyFieldToHeaderPlugin for model extraction and no response plugins.
 func NewBBRHarness(t *testing.T, ctx context.Context, streaming bool) *BBRHarness {
 	t.Helper()
 	modelToHeaderPlugin, err := bodyfieldtoheader.NewBodyFieldToHeaderPlugin(modelField, bodyfieldtoheader.ModelHeader)
@@ -94,11 +94,18 @@ func NewBBRHarness(t *testing.T, ctx context.Context, streaming bool) *BBRHarnes
 
 	baseModelToHeaderPlugin := &basemodelextractor.BaseModelToHeaderPlugin{AdaptersStore: store}
 
-	return NewBBRHarnessWithPlugins(t, ctx, streaming, []framework.RequestProcessor{modelToHeaderPlugin, baseModelToHeaderPlugin})
+	return NewBBRHarnessWithPlugins(t, ctx, streaming, []framework.RequestProcessor{modelToHeaderPlugin, baseModelToHeaderPlugin}, []framework.ResponseProcessor{})
 }
 
-// NewBBRHarnessWithPlugins boots up an isolated BBR server with custom request plugins.
-func NewBBRHarnessWithPlugins(t *testing.T, ctx context.Context, streaming bool, requestPlugins []framework.RequestProcessor) *BBRHarness {
+// NewBBRHarnessWithPlugins boots up an isolated BBR server on a random port
+// with the given request and response plugins.
+func NewBBRHarnessWithPlugins(
+	t *testing.T,
+	ctx context.Context,
+	streaming bool,
+	requestPlugins []framework.RequestProcessor,
+	responsePlugins []framework.ResponseProcessor,
+) *BBRHarness {
 	t.Helper()
 
 	// 1. Allocate Free Port
@@ -110,6 +117,7 @@ func NewBBRHarnessWithPlugins(t *testing.T, ctx context.Context, streaming bool,
 	runner.SecureServing = false
 	runner.Streaming = streaming
 	runner.RequestPlugins = requestPlugins
+	runner.ResponsePlugins = responsePlugins
 
 	// 3. Start Server in Background
 	serverCtx, serverCancel := context.WithCancel(ctx)

test/integration/bbr/hermetic_test.go

@@ -19,16 +19,22 @@ package bbr
 
 import (
 	"context"
+	"encoding/json"
 	"strings"
 	"testing"
 
+	envoyCorev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	envoyTypePb "github.com/envoyproxy/go-control-plane/envoy/type/v3"
 	"github.com/google/go-cmp/cmp"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/testing/protocmp"
 
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/framework"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/plugins/basemodelextractor"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/plugins/bodyfieldtoheader"
 	envoytest "sigs.k8s.io/gateway-api-inference-extension/pkg/common/envoy/test"
+	epp "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/plugin"
 	"sigs.k8s.io/gateway-api-inference-extension/test/integration"
 )
 
@@ -176,3 +182,175 @@ func TestFullDuplexStreamed_BodyBasedRouting(t *testing.T) {
 		})
 	}
 }
+
+// testResponsePlugin implements framework.ResponseProcessor for integration tests.
+type testResponsePlugin struct {
+	name     string
+	mutateFn func(ctx context.Context, response *framework.InferenceResponse) error
+}
+
+func (p *testResponsePlugin) TypedName() epp.TypedName {
+	return epp.TypedName{Type: "test", Name: p.name}
+}
+
+func (p *testResponsePlugin) ProcessResponse(ctx context.Context, _ *framework.CycleState, response *framework.InferenceResponse) error {
+	return p.mutateFn(ctx, response)
+}
+
+var _ framework.ResponseProcessor = &testResponsePlugin{}
+
+// TestResponsePlugins_Unary validates that response plugins can mutate the
+// response body in unary (non-streaming) mode over a real gRPC ext_proc stream.
+func TestResponsePlugins_Unary(t *testing.T) {
+	t.Parallel()
+
+	guardrailPlugin := &testResponsePlugin{
+		name: "guardrail",
+		mutateFn: func(_ context.Context, response *framework.InferenceResponse) error {
+			response.SetBodyField("guardrail", "applied")
+			return nil
+		},
+	}
+
+	ctx := context.Background()
+
+	modelToHeaderPlugin, err := bodyfieldtoheader.NewBodyFieldToHeaderPlugin(modelField, bodyfieldtoheader.ModelHeader)
+	require.NoError(t, err, "failed to create body-field-to-header plugin")
+	baseModelPlugin := &basemodelextractor.BaseModelToHeaderPlugin{AdaptersStore: basemodelextractor.NewAdaptersStore()}
+
+	h := NewBBRHarnessWithPlugins(t, ctx, false, []framework.RequestProcessor{modelToHeaderPlugin, baseModelPlugin}, []framework.ResponseProcessor{guardrailPlugin})
+
+	// Phase 1: Send request body (unary — no separate headers message).
+	requestBody := map[string]any{
+		"prompt":      "hello",
+		"max_tokens":  100,
+		"temperature": 0,
+		"model":       "test-model",
+	}
+	reqBodyBytes, err := json.Marshal(requestBody)
+	require.NoError(t, err)
+
+	// Phase 2: Build the full message sequence: RequestBody → ResponseHeaders → ResponseBody.
+	responseBody := map[string]any{
+		"choices": []any{
+			map[string]any{"text": "Hello!"},
+		},
+	}
+	respBodyBytes, err := json.Marshal(responseBody)
+	require.NoError(t, err)
+
+	reqs := []*extProcPb.ProcessingRequest{
+		{
+			Request: &extProcPb.ProcessingRequest_RequestBody{
+				RequestBody: &extProcPb.HttpBody{Body: reqBodyBytes, EndOfStream: true},
+			},
+		},
+		{
+			Request: &extProcPb.ProcessingRequest_ResponseHeaders{
+				ResponseHeaders: &extProcPb.HttpHeaders{
+					Headers: &envoyCorev3.HeaderMap{
+						Headers: []*envoyCorev3.HeaderValue{
+							{Key: "content-type", Value: "application/json"},
+						},
+					},
+				},
+			},
+		},
+		{
+			Request: &extProcPb.ProcessingRequest_ResponseBody{
+				ResponseBody: &extProcPb.HttpBody{Body: respBodyBytes, EndOfStream: true},
+			},
+		},
+	}
+
+	// Expect 3 responses: request body, response headers, response body.
+	responses, err := integration.StreamedRequest(t, h.Client, reqs, 3)
+	require.NoError(t, err, "unexpected error during streamed request")
+	require.Len(t, responses, 3)
+
+	// The response body should contain the guardrail field injected by the plugin.
+	expectedRespBody := map[string]any{
+		"choices": []any{
+			map[string]any{"text": "Hello!"},
+		},
+		"guardrail": "applied",
+	}
+
+	wantResponses := []*extProcPb.ProcessingResponse{
+		ExpectBBRUnaryResponse("test-model", "", "hello"),
+		ExpectResponseHeadersPassThrough(),
+		ExpectResponseBodyMutation(expectedRespBody),
+	}
+
+	envoytest.SortSetHeadersInResponses(wantResponses)
+	envoytest.SortSetHeadersInResponses(responses)
+	if diff := cmp.Diff(wantResponses, responses, protocmp.Transform()); diff != "" {
+		t.Errorf("Response mismatch (-want +got): %v", diff)
+	}
+}
+
+// TestResponsePlugins_NoPlugins_Unary validates that when no response plugins
+// are configured, the response body is passed through without mutation.
+func TestResponsePlugins_NoPlugins_Unary(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	h := NewBBRHarness(t, ctx, false)
+
+	requestBody := map[string]any{
+		"prompt":      "hello",
+		"max_tokens":  100,
+		"temperature": 0,
+		"model":       "test-model",
+	}
+	reqBodyBytes, err := json.Marshal(requestBody)
+	require.NoError(t, err)
+
+	responseBody := map[string]any{
+		"choices": []any{
+			map[string]any{"text": "Hi there!"},
+		},
+	}
+	respBodyBytes, err := json.Marshal(responseBody)
+	require.NoError(t, err)
+
+	reqs := []*extProcPb.ProcessingRequest{
+		{
+			Request: &extProcPb.ProcessingRequest_RequestBody{
+				RequestBody: &extProcPb.HttpBody{Body: reqBodyBytes, EndOfStream: true},
+			},
+		},
+		{
+			Request: &extProcPb.ProcessingRequest_ResponseHeaders{
+				ResponseHeaders: &extProcPb.HttpHeaders{
+					Headers: &envoyCorev3.HeaderMap{
+						Headers: []*envoyCorev3.HeaderValue{
+							{Key: "content-type", Value: "application/json"},
+						},
+					},
+				},
+			},
+		},
+		{
+			Request: &extProcPb.ProcessingRequest_ResponseBody{
+				ResponseBody: &extProcPb.HttpBody{Body: respBodyBytes, EndOfStream: true},
+			},
+		},
+	}
+
+	responses, err := integration.StreamedRequest(t, h.Client, reqs, 3)
+	require.NoError(t, err, "unexpected error during streamed request")
+	require.Len(t, responses, 3)
+
+	wantResponses := []*extProcPb.ProcessingResponse{
+		ExpectBBRUnaryResponse("test-model", "", "hello"),
+		ExpectResponseHeadersPassThrough(),
+		ExpectResponseBodyPassThrough(),
+	}
+
+	envoytest.SortSetHeadersInResponses(wantResponses)
+	envoytest.SortSetHeadersInResponses(responses)
+	if diff := cmp.Diff(wantResponses, responses, protocmp.Transform()); diff != "" {
+		t.Errorf("Response mismatch (-want +got): %v", diff)
+	}
+}

test/integration/bbr/util.go

@@ -18,6 +18,7 @@ package bbr
 
 import (
 	"encoding/json"
+	"strconv"
 
 	envoyCorev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
@@ -99,7 +100,58 @@ func ExpectBBRNoOpHeader() *extProcPb.ProcessingResponse {
 	}
 }
 
-// --- Response Expectations (Unary) ---
+// --- Response Phase Expectations ---
+
+// ExpectResponseHeadersPassThrough asserts that BBR passed response headers through with no mutations.
+func ExpectResponseHeadersPassThrough() *extProcPb.ProcessingResponse {
+	return &extProcPb.ProcessingResponse{
+		Response: &extProcPb.ProcessingResponse_ResponseHeaders{
+			ResponseHeaders: &extProcPb.HeadersResponse{},
+		},
+	}
+}
+
+// ExpectResponseBodyPassThrough asserts that BBR passed the response body through with no mutations
+// (i.e., no response plugins configured).
+func ExpectResponseBodyPassThrough() *extProcPb.ProcessingResponse {
+	return &extProcPb.ProcessingResponse{
+		Response: &extProcPb.ProcessingResponse_ResponseBody{
+			ResponseBody: &extProcPb.BodyResponse{},
+		},
+	}
+}
+
+// ExpectResponseBodyMutation asserts that a response plugin mutated the response body (unary mode).
+// Includes the Content-Length header mutation.
+func ExpectResponseBodyMutation(body map[string]any) *extProcPb.ProcessingResponse {
+	b, _ := json.Marshal(body)
+	return &extProcPb.ProcessingResponse{
+		Response: &extProcPb.ProcessingResponse_ResponseBody{
+			ResponseBody: &extProcPb.BodyResponse{
+				Response: &extProcPb.CommonResponse{
+					ClearRouteCache: true,
+					HeaderMutation: &extProcPb.HeaderMutation{
+						SetHeaders: []*envoyCorev3.HeaderValueOption{
+							{
+								Header: &envoyCorev3.HeaderValue{
+									Key:      "Content-Length",
+									RawValue: []byte(strconv.Itoa(len(b))),
+								},
+							},
+						},
+					},
+					BodyMutation: &extProcPb.BodyMutation{
+						Mutation: &extProcPb.BodyMutation_Body{
+							Body: b,
+						},
+					},
+				},
+			},
+		},
+	}
+}
+
+// --- Request Phase Expectations (Unary) ---
 
 // ExpectBBRUnaryResponse creates expected response for unary tests where the body is mutated directly.
 // baseModelName is the expected base model name (e.g., "llama" for both "llama" and "sql-lora-sheddable")