TF2OpenAPI E2E Tests (kubeflow#247)

* Add E2E tests for CLI

* Add JSON extension to golden files for highlighting

* Abstract out JSON expect statements for reuse

* Templatize tests

tools/tf2openapi/Makefile

@@ -5,6 +5,8 @@ TF_PROTO_OUT := generated
 # Run tests
 test: generate
 	go test ./types/... ./generator/...
+	go build -o bin/tf2openapi ./cmd/main.go
+	go test ./cmd/...
 
 # Generate code
 generate:

tools/tf2openapi/cmd/e2e_test.go

@@ -0,0 +1,182 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"testing"
+
+	"github.com/getkin/kin-openapi/openapi3"
+	"github.com/onsi/gomega"
+	gomegaTypes "github.com/onsi/gomega/types"
+
+	"github.com/kubeflow/kfserving/tools/tf2openapi/types"
+)
+
+// Functional E2E example
+func TestFlowers(t *testing.T) {
+	// model src: gs://kfserving-samples/models/tensorflow/flowers
+	g := gomega.NewGomegaWithT(t)
+	wd := workingDir(t)
+	cmdName := cmd(wd)
+	cmdArgs := []string{"--model_base_path", wd + "/testdata/" + t.Name() + ".pb"}
+	cmd := exec.Command(cmdName, cmdArgs...)
+	spec, err := cmd.Output()
+	g.Expect(err).Should(gomega.BeNil())
+	acceptableSpec := readFile("TestFlowers.golden.json", t)
+	acceptableSpecPermuted := readFile("TestFlowers2.golden.json", t)
+	g.Expect(spec).Should(gomega.Or(gomega.MatchJSON(acceptableSpec), gomega.MatchJSON(acceptableSpecPermuted)))
+}
+
+// Functional E2E example
+func TestCensusDifferentFlags(t *testing.T) {
+	// estimator model src: https://github.com/GoogleCloudPlatform/cloudml-samples/tree/master/census
+	g := gomega.NewGomegaWithT(t)
+	wd := workingDir(t)
+	cmdName := cmd(wd)
+	scenarios := map[string]struct {
+		cmdArgs      []string
+		expectedSpec []byte
+	}{
+		"Census": {
+			cmdArgs:      []string{"--model_base_path", wd + "/testdata/TestCensus.pb", "--signature_def", "predict"},
+			expectedSpec: readFile("TestCensus.golden.json", t),
+		},
+		"CustomFlags": {
+			cmdArgs: []string{"--model_base_path", wd + "/testdata/TestCustomFlags.pb",
+				"--signature_def", "predict", "--name", "customName", "--version", "1000",
+				"--metagraph_tags", "serve"},
+			expectedSpec: readFile("TestCustomFlags.golden.json", t),
+		},
+	}
+	for name, scenario := range scenarios {
+		t.Logf("Running %s ...", name)
+		cmd := exec.Command(cmdName, scenario.cmdArgs...)
+		spec, err := cmd.Output()
+		g.Expect(err).Should(gomega.BeNil())
+		swagger := &openapi3.Swagger{}
+		g.Expect(json.Unmarshal([]byte(spec), &swagger)).To(gomega.Succeed())
+
+		expectedSwagger := &openapi3.Swagger{}
+		g.Expect(json.Unmarshal(scenario.expectedSpec, &expectedSwagger)).To(gomega.Succeed())
+
+		// test equality, ignoring order in JSON arrays
+		expectJsonEquality(swagger, expectedSwagger, g)
+	}
+}
+
+func TestInputErrors(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+	wd := workingDir(t)
+	cmdName := cmd(wd)
+	scenarios := map[string]struct {
+		cmdArgs             []string
+		matchExpectedStdErr gomegaTypes.GomegaMatcher
+		expectedExitCode    int
+	}{
+		"InvalidCommand": {
+			cmdArgs:             []string{"--bad_flag"},
+			matchExpectedStdErr: gomega.And(gomega.ContainSubstring("Usage"), gomega.ContainSubstring("Flags:")),
+			expectedExitCode:    1,
+		},
+		"InvalidSavedModel": {
+			cmdArgs:             []string{"--model_base_path", wd + "/testdata/TestInvalidSavedModel.pb"},
+			matchExpectedStdErr: gomega.ContainSubstring(SavedModelFormatError),
+			expectedExitCode:    1,
+		},
+		"PropagateOpenAPIGenerationError": {
+			// model src: https://github.com/tensorflow/serving/tree/master/tensorflow_serving/example
+			cmdArgs: []string{"--model_base_path", wd + "/testdata/TestPropagateOpenAPIGenerationError.pb",
+				"--signature_def", "serving_default"},
+			matchExpectedStdErr: gomega.ContainSubstring(types.UnsupportedAPISchemaError),
+			expectedExitCode:    1,
+		},
+		"InvalidFilePath": {
+			cmdArgs:             []string{"--model_base_path", "badPath"},
+			matchExpectedStdErr: gomega.ContainSubstring(fmt.Sprintf(ModelBasePathError, "badPath", "")),
+			expectedExitCode:    1,
+		},
+	}
+	for name, scenario := range scenarios {
+		t.Logf("Running %s ...", name)
+		cmd := exec.Command(cmdName, scenario.cmdArgs...)
+		stdErr, err := cmd.CombinedOutput()
+		g.Expect(stdErr).Should(scenario.matchExpectedStdErr)
+		g.Expect(err.(*exec.ExitError).ExitCode()).To(gomega.Equal(scenario.expectedExitCode))
+	}
+}
+
+func TestOutputToFile(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+	wd := workingDir(t)
+	cmdName := cmd(wd)
+	defer os.Remove(wd + "/testdata/" + t.Name() + ".json")
+	outputFileName := t.Name() + ".json"
+	outputFilePath := wd + "/testdata/" + outputFileName
+	cmdArgs := []string{"--model_base_path", wd + "/testdata/TestFlowers.pb",
+		"--output_file", outputFilePath}
+	cmd := exec.Command(cmdName, cmdArgs...)
+	stdErr, err := cmd.CombinedOutput()
+	g.Expect(stdErr).To(gomega.BeEmpty())
+	g.Expect(err).Should(gomega.BeNil())
+	spec := readFile(outputFileName, t)
+	acceptableSpec := readFile("TestFlowers.golden.json", t)
+	acceptableSpecPermuted := readFile("TestFlowers2.golden.json", t)
+	g.Expect(spec).Should(gomega.Or(gomega.MatchJSON(acceptableSpec), gomega.MatchJSON(acceptableSpecPermuted)))
+}
+
+func TestOutputToFileTargetDirectoryError(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+	wd := workingDir(t)
+	cmdName := cmd(wd)
+	defer os.Remove(wd + "/testdata/" + t.Name() + ".json")
+	outputFileName := t.Name() + ".json"
+	badOutputFilePath := wd + "/nonexistent/" + outputFileName
+	cmdArgs := []string{"--model_base_path", wd + "/testdata/TestFlowers.pb",
+		"--output_file", badOutputFilePath}
+	cmd := exec.Command(cmdName, cmdArgs...)
+	stdErr, err := cmd.CombinedOutput()
+	g.Expect(err.(*exec.ExitError).ExitCode()).To(gomega.Equal(1))
+	expectedErr := fmt.Sprintf(OutputFilePathError, badOutputFilePath, "")
+	g.Expect(stdErr).Should(gomega.ContainSubstring(expectedErr))
+}
+
+func readFile(fName string, t *testing.T) []byte {
+	fPath := filepath.Join("testdata", fName)
+	openAPI, err := ioutil.ReadFile(fPath)
+	if err != nil {
+		t.Fatalf("failed reading %s: %s", fPath, err)
+	}
+	return openAPI
+}
+
+func workingDir(t *testing.T) string {
+	wd, err := os.Getwd()
+	if err != nil {
+		t.Fatalf("Failed os.Getwd() = %v, %v", wd, err)
+	}
+	return wd
+}
+
+func cmd(wd string) string {
+	return filepath.Dir(wd) + "/bin/tf2openapi"
+}
+
+func expectJsonEquality(swagger *openapi3.Swagger, expectedSwagger *openapi3.Swagger, g *gomega.GomegaWithT) {
+	instances := swagger.Components.RequestBodies["modelInput"].Value.Content.Get("application/json").
+		Schema.Value.Properties["instances"].Value
+	expectedInstances := expectedSwagger.Components.RequestBodies["modelInput"].Value.Content.
+		Get("application/json").Schema.Value.Properties["instances"].Value
+	g.Expect(swagger.Paths).Should(gomega.Equal(expectedSwagger.Paths))
+	g.Expect(swagger.OpenAPI).Should(gomega.Equal(expectedSwagger.OpenAPI))
+	g.Expect(swagger.Info).Should(gomega.Equal(expectedSwagger.Info))
+	g.Expect(swagger.Components.Responses).Should(gomega.Equal(expectedSwagger.Components.Responses))
+	g.Expect(instances.Items.Value.Required).Should(gomega.Not(gomega.BeNil()))
+	g.Expect(instances.Items.Value.Required).To(gomega.ConsistOf(expectedInstances.Items.Value.Required))
+	g.Expect(instances.Items.Value.AdditionalPropertiesAllowed).Should(gomega.Not(gomega.BeNil()))
+	g.Expect(instances.Items.Value.AdditionalPropertiesAllowed).Should(gomega.Equal(expectedInstances.Items.Value.AdditionalPropertiesAllowed))
+	g.Expect(instances.Items.Value.Properties).Should(gomega.Equal(expectedInstances.Items.Value.Properties))
+}

tools/tf2openapi/cmd/main.go

@@ -5,11 +5,19 @@ import (
 	pb "github.com/kubeflow/kfserving/tools/tf2openapi/generated/protobuf"
 	"github.com/kubeflow/kfserving/tools/tf2openapi/generator"
 	"github.com/spf13/cobra"
+	"fmt"
 	"io/ioutil"
 	"log"
 	"os"
 )
 
+// Known error messages
+const (
+	ModelBasePathError    = "Error reading file %s \n%s"
+	OutputFilePathError = "Failed writing to %s: %s"
+	SavedModelFormatError = "SavedModel not in expected format. May be corrupted: "
+)
+
 var (
 	modelBasePath string
 	modelName     string
@@ -46,9 +54,8 @@ func main() {
 func viewAPI(cmd *cobra.Command, args []string) {
 	modelPb, err := ioutil.ReadFile(modelBasePath)
 	if err != nil {
-		log.Fatalf("Error reading file %s \n%s", modelBasePath, err.Error())
+		log.Fatalf(ModelBasePathError, modelBasePath, err.Error())
 	}
-
 	generatorBuilder := &generator.Builder{}
 	if modelName != "" {
 		generatorBuilder.SetName(modelName)
@@ -63,34 +70,34 @@ func viewAPI(cmd *cobra.Command, args []string) {
 		generatorBuilder.SetMetaGraphTags(metaGraphTags)
 	}
 
-	model := UnmarshalSavedModelPb(modelPb)
+	model := unmarshalSavedModelPb(modelPb)
 	gen := generatorBuilder.Build()
-	spec, err :=  gen.GenerateOpenAPI(model)
+	spec, err := gen.GenerateOpenAPI(model)
 	if err != nil {
 		log.Fatalln(err.Error())
 	}
 	if outFile != "" {
 		f, err := os.Create(outFile)
 		if err != nil {
-			panic(err)
+			log.Fatalf(OutputFilePathError, outFile, err)
 		}
 		defer f.Close()
 		if _, err = f.WriteString(spec); err != nil {
 			panic(err)
 		}
 	} else {
 		// Default to std::out
-		log.Println(spec)
+		fmt.Println(spec)
 	}
 }
 
 /**
 Raises errors when model is missing fields that would pose an issue for Schema generation
- */
-func UnmarshalSavedModelPb(modelPb []byte) *pb.SavedModel {
+*/
+func unmarshalSavedModelPb(modelPb []byte) *pb.SavedModel {
 	model := &pb.SavedModel{}
 	if err := proto.Unmarshal(modelPb, model); err != nil {
-		log.Fatalln("SavedModel not in expected format. May be corrupted: " + err.Error())
+		log.Fatalln(SavedModelFormatError + err.Error())
 	}
 	return model
 }

tools/tf2openapi/cmd/testdata/TestCensus.golden.json

@@ -0,0 +1 @@
+{"components":{"requestBodies":{"modelInput":{"content":{"application/json":{"schema":{"additionalProperties":false,"properties":{"instances":{"items":{"additionalProperties":false,"properties":{"age":{"type":"number"},"capital_gain":{"type":"number"},"capital_loss":{"type":"number"},"education":{"type":"string"},"education_num":{"type":"number"},"gender":{"type":"string"},"hours_per_week":{"type":"number"},"marital_status":{"type":"string"},"native_country":{"type":"string"},"occupation":{"type":"string"},"race":{"type":"string"},"relationship":{"type":"string"},"workclass":{"type":"string"}},"required":["workclass","age","education","capital_loss","hours_per_week","relationship","marital_status","native_country","education_num","capital_gain","gender","occupation","race"],"type":"object"},"type":"array"}},"required":["instances"],"type":"object"}}}}},"responses":{"modelOutput":{"description":"Model output"}}},"info":{"title":"TFServing Predict Request API","version":"1.0"},"openapi":"3.0.0","paths":{"/v1/models/model/versions/1:predict":{"post":{"requestBody":{"$ref":"#/components/requestBodies/modelInput"},"responses":{"200":{"$ref":"#/components/responses/modelOutput"}}}}}}

tools/tf2openapi/cmd/testdata/TestCustomFlags.golden.json

@@ -0,0 +1 @@
+{"components":{"requestBodies":{"modelInput":{"content":{"application/json":{"schema":{"additionalProperties":false,"properties":{"instances":{"items":{"additionalProperties":false,"properties":{"age":{"type":"number"},"capital_gain":{"type":"number"},"capital_loss":{"type":"number"},"education":{"type":"string"},"education_num":{"type":"number"},"gender":{"type":"string"},"hours_per_week":{"type":"number"},"marital_status":{"type":"string"},"native_country":{"type":"string"},"occupation":{"type":"string"},"race":{"type":"string"},"relationship":{"type":"string"},"workclass":{"type":"string"}},"required":["workclass","age","education","capital_loss","hours_per_week","relationship","marital_status","native_country","education_num","capital_gain","gender","occupation","race"],"type":"object"},"type":"array"}},"required":["instances"],"type":"object"}}}}},"responses":{"modelOutput":{"description":"Model output"}}},"info":{"title":"TFServing Predict Request API","version":"1.0"},"openapi":"3.0.0","paths":{"/v1/models/customName/versions/1000:predict":{"post":{"requestBody":{"$ref":"#/components/requestBodies/modelInput"},"responses":{"200":{"$ref":"#/components/responses/modelOutput"}}}}}}

tools/tf2openapi/cmd/testdata/TestFlowers.golden.json

@@ -0,0 +1 @@
+{"components":{"requestBodies":{"modelInput":{"content":{"application/json":{"schema":{"additionalProperties":false,"properties":{"instances":{"items":{"additionalProperties":false,"properties":{"image_bytes":{"properties":{"b64":{"type":"string"}},"type":"object"},"key":{"type":"string"}},"required":["key","image_bytes"],"type":"object"},"type":"array"}},"required":["instances"],"type":"object"}}}}},"responses":{"modelOutput":{"description":"Model output"}}},"info":{"title":"TFServing Predict Request API","version":"1.0"},"openapi":"3.0.0","paths":{"/v1/models/model/versions/1:predict":{"post":{"requestBody":{"$ref":"#/components/requestBodies/modelInput"},"responses":{"200":{"$ref":"#/components/responses/modelOutput"}}}}}}

tools/tf2openapi/cmd/testdata/TestFlowers2.golden.json

@@ -0,0 +1 @@
+{"components":{"requestBodies":{"modelInput":{"content":{"application/json":{"schema":{"additionalProperties":false,"properties":{"instances":{"items":{"additionalProperties":false,"properties":{"image_bytes":{"properties":{"b64":{"type":"string"}},"type":"object"},"key":{"type":"string"}},"required":["image_bytes","key"],"type":"object"},"type":"array"}},"required":["instances"],"type":"object"}}}}},"responses":{"modelOutput":{"description":"Model output"}}},"info":{"title":"TFServing Predict Request API","version":"1.0"},"openapi":"3.0.0","paths":{"/v1/models/model/versions/1:predict":{"post":{"requestBody":{"$ref":"#/components/requestBodies/modelInput"},"responses":{"200":{"$ref":"#/components/responses/modelOutput"}}}}}}

tools/tf2openapi/cmd/testdata/TestInvalidSavedModel.pb

@@ -0,0 +1 @@
+this should not be unmarshalled