add initial integration test for body-based routing extension

Author: rramkumar1

test/integration/bbr/hermetic_test.go

@@ -0,0 +1,210 @@
+/*
+Copyright 2025 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+// Package bbr contains integration tests for the body-based routing extension.
+package bbr
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"testing"
+	"time"
+
+	configPb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
+	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+	"github.com/go-logr/logr"
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/protobuf/testing/protocmp"
+	"k8s.io/client-go/rest"
+	ctrl "sigs.k8s.io/controller-runtime"
+	runserver "sigs.k8s.io/gateway-api-inference-extension/pkg/body-based-routing/server"
+	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
+)
+
+const port = runserver.DefaultGrpcPort
+
+var (
+	serverRunner *runserver.ExtProcServerRunner
+	logger       = logutil.NewTestLogger().V(logutil.VERBOSE)
+)
+
+func TestBodyBasedRouting(t *testing.T) {
+	tests := []struct {
+		name              string
+		req               *extProcPb.ProcessingRequest
+		wantHeaders       []*configPb.HeaderValueOption
+		wantBody          []byte
+		wantErr           bool
+		immediateResponse *extProcPb.ImmediateResponse
+	}{
+		{
+			name: "success adding model parameter to header",
+			req:  generateRequest(logger, "llama"),
+			wantHeaders: []*configPb.HeaderValueOption{
+				{
+					Header: &configPb.HeaderValue{
+						Key:      "X-Gateway-Model-Name",
+						RawValue: []byte("llama"),
+					},
+				},
+				{
+					Header: &configPb.HeaderValue{
+						Key:      "Content-Length",
+						RawValue: []byte("76"),
+					},
+				},
+			},
+			wantBody: []byte("{\"max_tokens\":100,\"model\":\"llama\",\"prompt\":\"test1\",\"temperature\":0}"),
+			wantErr:  false,
+		},
+	}
+
+	// Set up extproc server runner with test environment config
+	cleanup := BeforeSuit(t)
+	defer cleanup()
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			client, cleanup := setUpHermeticServer()
+			t.Cleanup(cleanup)
+
+			want := &extProcPb.ProcessingResponse{
+				Response: &extProcPb.ProcessingResponse_RequestBody{
+					RequestBody: &extProcPb.BodyResponse{
+						Response: &extProcPb.CommonResponse{
+							HeaderMutation: &extProcPb.HeaderMutation{
+								SetHeaders: test.wantHeaders,
+							},
+							BodyMutation: &extProcPb.BodyMutation{
+								Mutation: &extProcPb.BodyMutation_Body{
+									Body: test.wantBody,
+								},
+							},
+						},
+					},
+				},
+			}
+
+			res, err := sendRequest(t, client, test.req)
+			if err != nil && !test.wantErr {
+				t.Errorf("Unexpected error, got: %v, want error: %v", err, test.wantErr)
+			}
+			if test.immediateResponse != nil {
+				want = &extProcPb.ProcessingResponse{
+					Response: &extProcPb.ProcessingResponse_ImmediateResponse{
+						ImmediateResponse: test.immediateResponse,
+					},
+				}
+			}
+			if diff := cmp.Diff(want, res, protocmp.Transform()); diff != "" {
+				t.Errorf("Unexpected response, (-want +got): %v", diff)
+			}
+		})
+	}
+}
+
+func generateRequest(logger logr.Logger, model string) *extProcPb.ProcessingRequest {
+	j := map[string]interface{}{
+		"prompt":      "test1",
+		"max_tokens":  100,
+		"temperature": 0,
+	}
+	if model != "" {
+		j["model"] = model
+	}
+
+	llmReq, err := json.Marshal(j)
+	if err != nil {
+		logutil.Fatal(logger, err, "Failed to unmarshal LLM request")
+	}
+	req := &extProcPb.ProcessingRequest{
+		Request: &extProcPb.ProcessingRequest_RequestBody{
+			RequestBody: &extProcPb.HttpBody{Body: llmReq},
+		},
+	}
+	return req
+}
+
+func setUpHermeticServer() (client extProcPb.ExternalProcessor_ProcessClient, cleanup func()) {
+	serverCtx, stopServer := context.WithCancel(context.Background())
+
+	go func() {
+		if err := serverRunner.AsRunnable(logger.WithName("ext-proc")).Start(serverCtx); err != nil {
+			logutil.Fatal(logger, err, "Failed to start ext-proc server")
+		}
+	}()
+
+	address := fmt.Sprintf("localhost:%v", port)
+	// Create a grpc connection
+	conn, err := grpc.NewClient(address, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		logutil.Fatal(logger, err, "Failed to connect", "address", address)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	client, err = extProcPb.NewExternalProcessorClient(conn).Process(ctx)
+	if err != nil {
+		logutil.Fatal(logger, err, "Failed to create client")
+	}
+	return client, func() {
+		cancel()
+		conn.Close()
+		stopServer()
+
+		// wait a little until the goroutines actually exit
+		time.Sleep(5 * time.Second)
+	}
+}
+
+// Sets up a test environment and returns the runner struct
+func BeforeSuit(t *testing.T) func() {
+	// Init runtime.
+	ctrl.SetLogger(logger)
+
+	mgr, err := ctrl.NewManager(&rest.Config{}, ctrl.Options{})
+	if err != nil {
+		logutil.Fatal(logger, err, "Failed to create manager")
+	}
+
+	serverRunner = runserver.NewDefaultExtProcServerRunner()
+	go func() {
+		if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
+			logutil.Fatal(logger, err, "Failed to start manager")
+		}
+	}()
+	logger.Info("Setting up hermetic ExtProc server")
+
+	return func() {}
+}
+
+func sendRequest(t *testing.T, client extProcPb.ExternalProcessor_ProcessClient, req *extProcPb.ProcessingRequest) (*extProcPb.ProcessingResponse, error) {
+	t.Logf("Sending request: %v", req)
+	if err := client.Send(req); err != nil {
+		t.Logf("Failed to send request %+v: %v", req, err)
+		return nil, err
+	}
+
+	res, err := client.Recv()
+	if err != nil {
+		t.Logf("Failed to receive: %v", err)
+		return nil, err
+	}
+	t.Logf("Received request %+v", res)
+	return res, err
+}