Skip to content

Commit 337abbf

Browse files
nirrozenbaumnirrozenbaum
authored
simplify bbr plugins and update tests accordingly (#2783)
* simplify bbr plugins and update tests accordingly Signed-off-by: Nir Rozenbaum <nrozenba@redhat.com> * code review comments Signed-off-by: Nir Rozenbaum <nrozenba@redhat.com> --------- Signed-off-by: Nir Rozenbaum <nrozenba@redhat.com>
1 parent 7eabe43 commit 337abbf

9 files changed

Lines changed: 35 additions & 170 deletions

File tree

‎config/charts/body-based-routing/values.yaml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,8 @@ bbr:
2222
- type: body-field-to-header
2323
name: model-extractor
2424
json:
25-
field_name: model
26-
header_name: X-Gateway-Model-Name
25+
fieldName: model
26+
headerName: X-Gateway-Model-Name
2727
- type: base-model-to-header
2828
name: base-model-mapper
2929

‎pkg/bbr/handlers/request_test.go‎

Lines changed: 1 addition & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -216,15 +216,7 @@ func TestHandleRequestBody(t *testing.T) {
216216
RequestBody: &extProcPb.BodyResponse{
217217
Response: &extProcPb.CommonResponse{
218218
ClearRouteCache: true,
219-
HeaderMutation: &extProcPb.HeaderMutation{
220-
SetHeaders: []*basepb.HeaderValueOption{
221-
{
222-
Header: &basepb.HeaderValue{
223-
Key: basemodelextractor.BaseModelHeader,
224-
},
225-
},
226-
},
227-
},
219+
HeaderMutation: &extProcPb.HeaderMutation{},
228220
},
229221
},
230222
},
@@ -249,11 +241,6 @@ func TestHandleRequestBody(t *testing.T) {
249241
RawValue: []byte("27"),
250242
},
251243
},
252-
{
253-
Header: &basepb.HeaderValue{
254-
Key: basemodelextractor.BaseModelHeader,
255-
},
256-
},
257244
},
258245
},
259246
},
@@ -655,12 +642,6 @@ func TestHandleRequestBody_BodyMutation(t *testing.T) {
655642
ClearRouteCache: true,
656643
HeaderMutation: &extProcPb.HeaderMutation{
657644
SetHeaders: []*basepb.HeaderValueOption{
658-
{
659-
Header: &basepb.HeaderValue{
660-
Key: basemodelextractor.BaseModelHeader,
661-
RawValue: []byte(""),
662-
},
663-
},
664645
{
665646
Header: &basepb.HeaderValue{
666647
Key: contentLengthHeader,
@@ -697,12 +678,6 @@ func TestHandleRequestBody_BodyMutation(t *testing.T) {
697678
ClearRouteCache: true,
698679
HeaderMutation: &extProcPb.HeaderMutation{
699680
SetHeaders: []*basepb.HeaderValueOption{
700-
{
701-
Header: &basepb.HeaderValue{
702-
Key: basemodelextractor.BaseModelHeader,
703-
RawValue: []byte(""),
704-
},
705-
},
706681
{
707682
Header: &basepb.HeaderValue{
708683
Key: contentLengthHeader,

‎pkg/bbr/plugins/basemodelextractor/base_model_to_header.go‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -86,16 +86,11 @@ func (p *BaseModelToHeaderPlugin) WithName(name string) *BaseModelToHeaderPlugin
8686

8787
// ProcessRequest sets base model name on the header
8888
func (p *BaseModelToHeaderPlugin) ProcessRequest(ctx context.Context, _ *framework.CycleState, request *framework.InferenceRequest) error {
89-
if request == nil || request.Headers == nil || request.Body == nil {
90-
return nil // this shouldn't happen
91-
}
92-
9389
// extract raw field value from body
9490
rawFieldValue, exists := request.Body[modelField]
9591
if !exists {
96-
// If model field is not present, set empty base model header
97-
request.SetHeader(BaseModelHeader, "")
98-
log.FromContext(ctx).V(logutil.VERBOSE).Info("model field not found, setting empty base model header")
92+
// If model field is not present, skip base model header
93+
log.FromContext(ctx).V(logutil.VERBOSE).Info("model field not found, skipping")
9994
return nil
10095
}
10196

‎pkg/bbr/plugins/basemodelextractor/base_model_to_header_test.go‎

Lines changed: 7 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -46,10 +46,10 @@ func TestBaseModelToHeaderPlugin_TypedName(t *testing.T) {
4646

4747
got := p.TypedName()
4848
if got.Type != BaseModelToHeaderPluginType {
49-
t.Errorf("TypedName().Type = %q, want %q", got.Type, BaseModelToHeaderPluginType)
49+
t.Errorf("TypedName().Type = %s, want %s", got.Type, BaseModelToHeaderPluginType)
5050
}
5151
if got.Name != "test-plugin" {
52-
t.Errorf("TypedName().Name = %q, want %q", got.Name, "test-plugin")
52+
t.Errorf("TypedName().Name = %s, want %s", got.Name, "test-plugin")
5353
}
5454
}
5555

@@ -64,10 +64,10 @@ func TestBaseModelToHeaderPlugin_WithName(t *testing.T) {
6464

6565
got := p.TypedName()
6666
if got.Name != "new-name" {
67-
t.Errorf("TypedName().Name = %q, want %q", got.Name, "new-name")
67+
t.Errorf("TypedName().Name = %s, want %s", got.Name, "new-name")
6868
}
6969
if got.Type != BaseModelToHeaderPluginType {
70-
t.Errorf("TypedName().Type = %q, want %q", got.Type, BaseModelToHeaderPluginType)
70+
t.Errorf("TypedName().Type = %s, want %s", got.Type, BaseModelToHeaderPluginType)
7171
}
7272
}
7373

@@ -77,7 +77,6 @@ func TestBaseModelToHeaderPluginFactory(t *testing.T) {
7777
name string
7878
pluginName string
7979
rawParams json.RawMessage
80-
wantErr bool
8180
wantName string
8281
}{
8382
{
@@ -121,20 +120,14 @@ func TestBaseModelToHeaderPluginFactory(t *testing.T) {
121120
handle := framework.NewBbrHandle(context.Background(), mgr)
122121

123122
p, err := BaseModelToHeaderPluginFactory(tt.pluginName, tt.rawParams, handle)
124-
if tt.wantErr {
125-
if err == nil {
126-
t.Fatal("expected error, got nil")
127-
}
128-
return
129-
}
130123
if err != nil {
131124
t.Fatalf("unexpected error: %v", err)
132125
}
133126
if got := p.TypedName().Name; got != tt.wantName {
134-
t.Errorf("Name = %q, want %q", got, tt.wantName)
127+
t.Errorf("Name = %s, want %s", got, tt.wantName)
135128
}
136129
if got := p.TypedName().Type; got != BaseModelToHeaderPluginType {
137-
t.Errorf("Type = %q, want %q", got, BaseModelToHeaderPluginType)
130+
t.Errorf("Type = %s, want %s", got, BaseModelToHeaderPluginType)
138131
}
139132
})
140133
}
@@ -242,28 +235,6 @@ func TestBaseModelToHeaderPlugin_ProcessRequest(t *testing.T) {
242235
}(),
243236
wantHeader: "",
244237
},
245-
{
246-
name: "nil request",
247-
request: nil,
248-
},
249-
{
250-
name: "nil headers",
251-
request: &framework.InferenceRequest{
252-
InferenceMessage: framework.InferenceMessage{
253-
Headers: nil,
254-
Body: map[string]any{"model": testBaseModel},
255-
},
256-
},
257-
},
258-
{
259-
name: "nil body",
260-
request: &framework.InferenceRequest{
261-
InferenceMessage: framework.InferenceMessage{
262-
Headers: map[string]string{},
263-
Body: nil,
264-
},
265-
},
266-
},
267238
}
268239
for _, tt := range tests {
269240
t.Run(tt.name, func(t *testing.T) {
@@ -279,7 +250,7 @@ func TestBaseModelToHeaderPlugin_ProcessRequest(t *testing.T) {
279250
}
280251
if tt.wantHeader != "" && tt.request != nil && tt.request.Headers != nil {
281252
if got := tt.request.Headers[BaseModelHeader]; got != tt.wantHeader {
282-
t.Errorf("Headers[%q] = %q, want %q", BaseModelHeader, got, tt.wantHeader)
253+
t.Errorf("Headers[%s] = %s, want %s", BaseModelHeader, got, tt.wantHeader)
283254
}
284255
}
285256
})

‎pkg/bbr/plugins/bodyfieldtoheader/body_field_to_header.go‎

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@ package bodyfieldtoheader
1919
import (
2020
"context"
2121
"encoding/json"
22-
"errors"
2322
"fmt"
2423

2524
"sigs.k8s.io/controller-runtime/pkg/log"
@@ -41,9 +40,9 @@ var _ framework.RequestProcessor = &BodyFieldToHeaderPlugin{}
4140
// BodyFieldToHeaderConfig defines the JSON configuration structure for the plugin.
4241
type BodyFieldToHeaderConfig struct {
4342
// FieldName is the name of the body field to extract
44-
FieldName string `json:"field_name"`
43+
FieldName string `json:"fieldName"`
4544
// HeaderName is the name of the header to set
46-
HeaderName string `json:"header_name"`
45+
HeaderName string `json:"headerName"`
4746
}
4847

4948
// BodyFieldToHeaderPluginFactory defines the factory function for NewBodyFieldToHeaderPlugin.
@@ -67,11 +66,11 @@ func BodyFieldToHeaderPluginFactory(name string, rawParameters json.RawMessage,
6766
// NewBodyFieldToHeaderPlugin initializes a new BodyFieldToHeaderPlugin and returns its pointer.
6867
func NewBodyFieldToHeaderPlugin(fieldName, headerName string) (*BodyFieldToHeaderPlugin, error) {
6968
if fieldName == "" {
70-
return nil, errors.New("body fieldName is required in BodyFieldToHeader plugin")
69+
return nil, fmt.Errorf("fieldName is required for plugin '%s'", BodyFieldToHeaderPluginType)
7170
}
7271

7372
if headerName == "" {
74-
return nil, errors.New("headerName is required in BodyFieldToHeader plugin")
73+
return nil, fmt.Errorf("headerName is required for plugin '%s'", BodyFieldToHeaderPluginType)
7574
}
7675

7776
return &BodyFieldToHeaderPlugin{
@@ -104,10 +103,6 @@ func (p *BodyFieldToHeaderPlugin) WithName(name string) *BodyFieldToHeaderPlugin
104103

105104
// ProcessRequest extracts value from a given body field and sets it as HTTP header.
106105
func (p *BodyFieldToHeaderPlugin) ProcessRequest(ctx context.Context, _ *framework.CycleState, request *framework.InferenceRequest) error {
107-
if request == nil || request.Headers == nil || request.Body == nil {
108-
return nil // this shouldn't happen
109-
}
110-
111106
// extract raw field value from body
112107
rawFieldValue, exists := request.Body[p.fieldName]
113108
if !exists {

‎pkg/bbr/plugins/bodyfieldtoheader/body_field_to_header_test.go‎

Lines changed: 14 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,9 @@ import (
2424
"sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/framework"
2525
)
2626

27-
const testModelValue = "qwen3"
27+
const (
28+
testModelValue = "qwen3"
29+
)
2830

2931
func TestNewBodyFieldToHeaderPlugin(t *testing.T) {
3032
tests := []struct {
@@ -70,10 +72,10 @@ func TestNewBodyFieldToHeaderPlugin(t *testing.T) {
7072
t.Fatalf("unexpected error: %v", err)
7173
}
7274
if p.TypedName().Type != BodyFieldToHeaderPluginType {
73-
t.Errorf("Type = %q, want %q", p.TypedName().Type, BodyFieldToHeaderPluginType)
75+
t.Errorf("Type = %s, want %s", p.TypedName().Type, BodyFieldToHeaderPluginType)
7476
}
7577
if p.TypedName().Name != BodyFieldToHeaderPluginType {
76-
t.Errorf("Name = %q, want %q", p.TypedName().Name, BodyFieldToHeaderPluginType)
78+
t.Errorf("Name = %s, want %s", p.TypedName().Name, BodyFieldToHeaderPluginType)
7779
}
7880
})
7981
}
@@ -88,10 +90,10 @@ func TestBodyFieldToHeaderPlugin_WithName(t *testing.T) {
8890
p.WithName("custom-name")
8991

9092
if got := p.TypedName().Name; got != "custom-name" {
91-
t.Errorf("Name after WithName = %q, want %q", got, "custom-name")
93+
t.Errorf("Name after WithName = %s, want %s", got, "custom-name")
9294
}
9395
if got := p.TypedName().Type; got != BodyFieldToHeaderPluginType {
94-
t.Errorf("Type should be unchanged = %q, want %q", got, BodyFieldToHeaderPluginType)
96+
t.Errorf("Type should be unchanged = %s, want %s", got, BodyFieldToHeaderPluginType)
9597
}
9698
}
9799

@@ -106,7 +108,7 @@ func TestBodyFieldToHeaderPluginFactory(t *testing.T) {
106108
{
107109
name: "valid config",
108110
pluginName: "my-plugin",
109-
rawParams: json.RawMessage(`{"field_name":"model","header_name":"X-Gateway-Model"}`),
111+
rawParams: json.RawMessage(`{"fieldName":"model","headerName":"X-Gateway-Model"}`),
110112
wantName: "my-plugin",
111113
},
112114
{
@@ -116,13 +118,13 @@ func TestBodyFieldToHeaderPluginFactory(t *testing.T) {
116118
wantErr: true,
117119
},
118120
{
119-
name: "missing field_name",
121+
name: "missing fieldName",
120122
pluginName: "my-plugin",
121123
rawParams: json.RawMessage(`{"header_name":"X-Gateway-Model"}`),
122124
wantErr: true,
123125
},
124126
{
125-
name: "missing header_name",
127+
name: "missing headerName",
126128
pluginName: "my-plugin",
127129
rawParams: json.RawMessage(`{"field_name":"model"}`),
128130
wantErr: true,
@@ -165,10 +167,10 @@ func TestBodyFieldToHeaderPluginFactory(t *testing.T) {
165167
t.Fatalf("unexpected error: %v", err)
166168
}
167169
if got := p.TypedName().Name; got != tt.wantName {
168-
t.Errorf("Name = %q, want %q", got, tt.wantName)
170+
t.Errorf("Name = %s, want %s", got, tt.wantName)
169171
}
170172
if got := p.TypedName().Type; got != BodyFieldToHeaderPluginType {
171-
t.Errorf("Type = %q, want %q", got, BodyFieldToHeaderPluginType)
173+
t.Errorf("Type = %s, want %s", got, BodyFieldToHeaderPluginType)
172174
}
173175
})
174176
}
@@ -247,45 +249,6 @@ func TestBodyFieldToHeaderPlugin_ProcessRequest(t *testing.T) {
247249
return r
248250
}(),
249251
},
250-
{
251-
name: "nil field value",
252-
fieldName: "model",
253-
headerName: "X-Gateway-Model",
254-
request: func() *framework.InferenceRequest {
255-
r := framework.NewInferenceRequest()
256-
r.Body["model"] = nil
257-
return r
258-
}(),
259-
wantHeader: "<nil>",
260-
},
261-
{
262-
name: "nil request",
263-
fieldName: "model",
264-
headerName: "X-Gateway-Model",
265-
request: nil,
266-
},
267-
{
268-
name: "nil headers",
269-
fieldName: "model",
270-
headerName: "X-Gateway-Model",
271-
request: &framework.InferenceRequest{
272-
InferenceMessage: framework.InferenceMessage{
273-
Headers: nil,
274-
Body: map[string]any{"model": "qwen"},
275-
},
276-
},
277-
},
278-
{
279-
name: "nil body",
280-
fieldName: "model",
281-
headerName: "X-Gateway-Model",
282-
request: &framework.InferenceRequest{
283-
InferenceMessage: framework.InferenceMessage{
284-
Headers: map[string]string{},
285-
Body: nil,
286-
},
287-
},
288-
},
289252
}
290253
for _, tt := range tests {
291254
t.Run(tt.name, func(t *testing.T) {
@@ -306,7 +269,7 @@ func TestBodyFieldToHeaderPlugin_ProcessRequest(t *testing.T) {
306269
}
307270
if tt.wantHeader != "" {
308271
if got := tt.request.Headers[tt.headerName]; got != tt.wantHeader {
309-
t.Errorf("Headers[%q] = %q, want %q", tt.headerName, got, tt.wantHeader)
272+
t.Errorf("Headers[%s] = %s, want %s", tt.headerName, got, tt.wantHeader)
310273
}
311274
}
312275
})
@@ -328,6 +291,6 @@ func TestBodyFieldToHeaderPlugin_ProcessRequest_MutatedHeaders(t *testing.T) {
328291

329292
mutated := request.MutatedHeaders()
330293
if got, ok := mutated["X-Gateway-Model"]; !ok || got != testModelValue {
331-
t.Errorf("MutatedHeaders[\"X-Gateway-Model\"] = %q, %v; want %q, true", got, ok, testModelValue)
294+
t.Errorf("MutatedHeaders[\"X-Gateway-Model\"] = %s, %v; want %s, true", got, ok, testModelValue)
332295
}
333296
}

‎test/integration/bbr/body_mutation_test.go‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -88,12 +88,6 @@ func TestBodyMutation_Unary(t *testing.T) {
8888
ClearRouteCache: true,
8989
HeaderMutation: &extProcPb.HeaderMutation{
9090
SetHeaders: []*envoyCorev3.HeaderValueOption{
91-
{
92-
Header: &envoyCorev3.HeaderValue{
93-
Key: "X-Gateway-Base-Model-Name",
94-
RawValue: []byte(""),
95-
},
96-
},
9791
{
9892
Header: &envoyCorev3.HeaderValue{
9993
Key: "Content-Length",
@@ -166,12 +160,6 @@ func TestBodyMutation_Streaming(t *testing.T) {
166160
ClearRouteCache: true,
167161
HeaderMutation: &extProcPb.HeaderMutation{
168162
SetHeaders: []*envoyCorev3.HeaderValueOption{
169-
{
170-
Header: &envoyCorev3.HeaderValue{
171-
Key: "X-Gateway-Base-Model-Name",
172-
RawValue: []byte(""),
173-
},
174-
},
175163
{
176164
Header: &envoyCorev3.HeaderValue{
177165
Key: "Content-Length",

0 commit comments

Comments
 (0)