Review updates

SeldonIO · Mar 25, 2020 · ee1874c · ee1874c
1 parent 72bdf77
commit ee1874c
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/executor/api/grpc/seldon/server_test.go b/executor/api/grpc/seldon/server_test.go
@@ -28,7 +28,7 @@ func TestPredict(t *testing.T) {
 		},
 	}
 	url, _ := url.Parse("http://localhost")
-	server := NewGrpcSeldonServer(&p, test.SeldonMessageTestClient{}, url, "default")
+	server := NewGrpcSeldonServer(&p, &test.SeldonMessageTestClient{}, url, "default")
 
 	var sm proto.SeldonMessage
 	var data = ` {"data":{"ndarray":[[1.1,2.0]]}}`

diff --git a/executor/predictor/predictor_process_test.go b/executor/predictor/predictor_process_test.go
@@ -43,10 +43,10 @@ func createPredictorProcessWithRoute(t *testing.T, chosenRoute int) *PredictorPr
 	return &pp
 }
 
-func createPredictorProcessWithError(t *testing.T, errMethod *v1.PredictiveUnitMethod, err error) *PredictorProcess {
+func createPredictorProcessWithError(t *testing.T, errMethod *v1.PredictiveUnitMethod, err error, errPayload payload.SeldonPayload) *PredictorProcess {
 	url, _ := url.Parse(testSourceUrl)
 	ctx := context.WithValue(context.TODO(), payload.SeldonPUIDHeader, testSeldonPuid)
-	pp := NewPredictorProcess(ctx, &test.SeldonMessageTestClient{ErrMethod: errMethod, Err: err}, logf.Log.WithName("SeldonMessageRestClient"), url, "default", map[string][]string{})
+	pp := NewPredictorProcess(ctx, &test.SeldonMessageTestClient{ErrMethod: errMethod, Err: err, ErrPayload: errPayload}, logf.Log.WithName("SeldonMessageRestClient"), url, "default", map[string][]string{})
 	return &pp
 }
 
@@ -320,9 +320,14 @@ func TestModelError(t *testing.T) {
 
 	errMethod := v1.TRANSFORM_INPUT
 	chosenErr := errors.New("something bad happened")
-	pResp, err := createPredictorProcessWithError(t, &errMethod, chosenErr).Predict(graph, createPredictPayload(g))
+	errBytes := "{\"status\":\"failed\"}"
+	errPayload := payload.BytesPayload{
+		Msg:         []byte(errBytes),
+		ContentType: "xyz",
+	}
+	pResp, err := createPredictorProcessWithError(t, &errMethod, chosenErr, &errPayload).Predict(graph, createPredictPayload(g))
 	g.Expect(err).ShouldNot(BeNil())
-	g.Expect(pResp).Should(BeNil())
+	g.Expect(pResp).To(Equal(&errPayload))
 	g.Expect(err.Error()).Should(Equal("something bad happened"))
 }