Added test for testing response.

go.mod

@@ -3,6 +3,7 @@ module github.com/kserve/rest-proxy
 go 1.16
 
 require (
+	github.com/google/go-cmp v0.5.6
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.6.0
 	google.golang.org/grpc v1.40.0
 	google.golang.org/protobuf v1.27.1

proxy/marshaler.go

@@ -82,7 +82,6 @@ var tensorTypes = map[string]tensorType{
 // returning to the user.
 func (c *CustomJSONPb) Marshal(v interface{}) ([]byte, error) {
 	if r, ok := v.(*gw.ModelInferResponse); ok {
-		fmt.Printf("%v, %v\n", v, r)
 		resp := &RESTResponse{}
 		resp.ModelName = r.ModelName
 		resp.ModelVersion = r.ModelVersion

proxy/marshaler_test.go

@@ -3,6 +3,7 @@ package main
 import (
 	"bytes"
 	"fmt"
+	"github.com/google/go-cmp/cmp"
 	"reflect"
 	"strings"
 	"testing"
@@ -93,6 +94,25 @@ var data4D = `
 	]
 `
 
+func generateProtoBufResponse() gw.ModelInferResponse {
+	expectedOutput := []*gw.ModelInferResponse_InferOutputTensor{{
+		Name:     "predict",
+		Datatype: "INT64",
+		Shape:    []int64{2},
+		Contents: &gw.InferTensorContents{
+			Int64Contents: []int64{8, 8},
+		},
+	}}
+
+	return gw.ModelInferResponse{
+		ModelName: "example",
+		Id:        "foo",
+		Outputs:   expectedOutput,
+	}
+}
+
+var jsonResponse = `{"model_name":"example","id":"foo","outputs":[{"name":"predict","datatype":"INT64","shape":[2],"contents":{"int64_contents":[8,8]}}]}`
+
 func generateProtoBufRequest(shape []int64) *gw.ModelInferRequest {
 	var expectedInput = gw.ModelInferRequest_InferInputTensor{
 		Name:     "predict",
@@ -118,7 +138,7 @@ func generateProtoBufRequest(shape []int64) *gw.ModelInferRequest {
 	return modelInferRequest
 }
 
-func TestRestReq(t *testing.T) {
+func TestRESTRequest(t *testing.T) {
 	c := CustomJSONPb{}
 	inputDataArray := []string{data1D, data2D, data3D, data4D}
 	inputDataShapes := [][]int64{{2, 64}, {2, 64}, {2, 2, 32}, {2, 2, 2, 16}}
@@ -136,4 +156,14 @@ func TestRestReq(t *testing.T) {
 	}
 }
 
-
+func TestRESTResponse(t *testing.T) {
+	c := CustomJSONPb{}
+	v := generateProtoBufResponse()
+	marshal, err := c.Marshal(v)
+	if err != nil {
+		t.Error(err)
+	}
+	if d := cmp.Diff(string(marshal), jsonResponse); d != "" {
+		t.Errorf("diff :%s", d)
+	}
+}