From b182d0d650678d008e766c01ea49850cfa60d95d Mon Sep 17 00:00:00 2001
From: Prashant Sharma <prashsh1@in.ibm.com>
Date: Thu, 2 Dec 2021 19:49:03 +0530
Subject: [PATCH] Added test for testing response.

---
 proxy/marshaler.go      |  1 -
 proxy/marshaler_test.go | 34 ++++++++++++++++++++++++++++++++--
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/proxy/marshaler.go b/proxy/marshaler.go
index 2a66add..442f828 100644
--- a/proxy/marshaler.go
+++ b/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
diff --git a/proxy/marshaler_test.go b/proxy/marshaler_test.go
index 5beea8a..63a7321 100644
--- a/proxy/marshaler_test.go
+++ b/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)
+	}
+}