Skip to content

Commit b62ad17

Browse files
committed
Add unit tests for request body
1 parent 3952e22 commit b62ad17

File tree

1 file changed

+9
-17
lines changed

1 file changed

+9
-17
lines changed

pkg/epp/handlers/request_test.go

+9-17
Original file line numberDiff line numberDiff line change
@@ -49,15 +49,15 @@ func TestHandleRequestBody(t *testing.T) {
4949
// Setup datastore
5050
tsModel := "food-review"
5151
modelWithTarget := "food-review-0"
52-
model1ts := testutil.MakeInferenceModel("model1").
52+
model1 := testutil.MakeInferenceModel("model1").
5353
CreationTimestamp(metav1.Unix(1000, 0)).
5454
ModelName(tsModel).ObjRef()
5555
model2 := testutil.MakeInferenceModel("model2").
5656
CreationTimestamp(metav1.Unix(1000, 0)).
5757
ModelName(modelWithTarget).ObjRef()
5858
pmf := backendmetrics.NewPodMetricsFactory(&backendmetrics.FakePodMetricsClient{}, time.Second)
5959
ds := datastore.NewDatastore(t.Context(), pmf)
60-
ds.ModelSetIfOlder(model1ts)
60+
ds.ModelSetIfOlder(model1)
6161
ds.ModelSetIfOlder(model2)
6262

6363
pool := &v1alpha2.InferencePool{
@@ -82,9 +82,7 @@ func TestHandleRequestBody(t *testing.T) {
8282
tests := []struct {
8383
name string
8484
reqBodyMap map[string]interface{}
85-
reqCtx *RequestContext
8685
wantErrCode string
87-
wantErr bool
8886
wantReqCtx *RequestContext
8987
wantRespBody map[string]interface{}
9088
}{
@@ -124,7 +122,6 @@ func TestHandleRequestBody(t *testing.T) {
124122
},
125123
{
126124
name: "no model defined, expect err",
127-
wantErr: true,
128125
wantErrCode: errutil.BadRequest,
129126
},
130127
{
@@ -133,7 +130,6 @@ func TestHandleRequestBody(t *testing.T) {
133130
"model": "non-existent-model",
134131
"prompt": "test prompt",
135132
},
136-
wantErr: true,
137133
wantErrCode: errutil.BadConfiguration,
138134
},
139135
{
@@ -142,22 +138,18 @@ func TestHandleRequestBody(t *testing.T) {
142138
"model": "food-review-1",
143139
"prompt": "test prompt",
144140
},
145-
wantErr: true,
146141
wantErrCode: errutil.BadConfiguration,
147142
},
148143
}
149144

150145
for _, test := range tests {
151146
t.Run(test.name, func(t *testing.T) {
152147
server := NewStreamingServer(scheduling.NewScheduler(ds), DefaultDestinationEndpointHintMetadataNamespace, DefaultDestinationEndpointHintKey, ds)
153-
reqCtx := test.reqCtx
154-
if reqCtx == nil {
155-
reqCtx = &RequestContext{}
156-
}
148+
reqCtx := &RequestContext{}
157149
req := &extProcPb.ProcessingRequest{}
158-
gotReqCtx, err := server.HandleRequestBody(ctx, reqCtx, req, test.reqBodyMap)
150+
reqCtx, err := server.HandleRequestBody(ctx, reqCtx, req, test.reqBodyMap)
159151

160-
if test.wantErr {
152+
if test.wantErrCode != "" {
161153
if err == nil {
162154
t.Fatalf("HandleRequestBody should have returned an error containing '%s', but got nil", test.wantErrCode)
163155
}
@@ -172,16 +164,16 @@ func TestHandleRequestBody(t *testing.T) {
172164
}
173165

174166
if test.wantReqCtx != nil {
175-
if diff := cmp.Diff(test.wantReqCtx.Model, gotReqCtx.Model); diff != "" {
167+
if diff := cmp.Diff(test.wantReqCtx.Model, reqCtx.Model); diff != "" {
176168
t.Errorf("HandleRequestBody returned unexpected reqCtx.Model, diff(-want, +got): %v", diff)
177169
}
178-
if diff := cmp.Diff(test.wantReqCtx.ResolvedTargetModel, gotReqCtx.ResolvedTargetModel); diff != "" {
170+
if diff := cmp.Diff(test.wantReqCtx.ResolvedTargetModel, reqCtx.ResolvedTargetModel); diff != "" {
179171
t.Errorf("HandleRequestBody returned unexpected reqCtx.ResolvedTargetModel, diff(-want, +got): %v", diff)
180172
}
181-
if diff := cmp.Diff(test.wantReqCtx.TargetPod, gotReqCtx.TargetPod); diff != "" {
173+
if diff := cmp.Diff(test.wantReqCtx.TargetPod, reqCtx.TargetPod); diff != "" {
182174
t.Errorf("HandleRequestBody returned unexpected reqCtx.TargetPod, diff(-want, +got): %v", diff)
183175
}
184-
if diff := cmp.Diff(test.wantReqCtx.TargetEndpoint, gotReqCtx.TargetEndpoint); diff != "" {
176+
if diff := cmp.Diff(test.wantReqCtx.TargetEndpoint, reqCtx.TargetEndpoint); diff != "" {
185177
t.Errorf("HandleRequestBody returned unexpected reqCtx.TargetEndpoint, diff(-want, +got): %v", diff)
186178
}
187179
}

0 commit comments

Comments
 (0)