From 5e474de1e9c13280a0d710256fb676e2413d7c09 Mon Sep 17 00:00:00 2001
From: Zaid Ajaj <zaid.naom@gmail.com>
Date: Wed, 14 Dec 2022 23:21:22 +0100
Subject: [PATCH] Add a yaml test schema

---
 pkg/codegen/schema/schema.go                  | 18 +++++++--------
 pkg/codegen/schema/schema_test.go             | 22 +++++++++++++++++--
 .../schema/good-simplified-methods.yml        | 14 ++++++++++++
 3 files changed, 43 insertions(+), 11 deletions(-)
 create mode 100644 pkg/codegen/testing/test/testdata/schema/good-simplified-methods.yml

diff --git a/pkg/codegen/schema/schema.go b/pkg/codegen/schema/schema.go
index ee1bc03dd06e..d72db6383f74 100644
--- a/pkg/codegen/schema/schema.go
+++ b/pkg/codegen/schema/schema.go
@@ -1579,9 +1579,9 @@ func (returnTypeSpec *ReturnTypeSpec) UnmarshalJSON(inputJSON []byte) error {
 	return nil
 }
 
-func (returnTypeSpec *ReturnTypeSpec) UnmarshalYAML(node *yaml.Node) error {
+func (returnTypeSpec *ReturnTypeSpec) UnmarshalYAML(inputData []byte) error {
 	var data map[string]interface{}
-	if err := node.Decode(&data); err != nil {
+	if err := yaml.Unmarshal(inputData, &data); err != nil {
 		return err
 	}
 
@@ -1592,11 +1592,11 @@ func (returnTypeSpec *ReturnTypeSpec) UnmarshalYAML(node *yaml.Node) error {
 	var objectSpec *ObjectTypeSpec
 	var typeSpec *TypeSpec
 	if _, hasProperties := data["properties"]; hasProperties {
-		if err := node.Decode(&objectSpec); err != nil {
+		if err := yaml.Unmarshal(inputData, &objectSpec); err != nil {
 			return err
 		}
 	} else {
-		if err := node.Decode(&typeSpec); err != nil {
+		if err := yaml.Unmarshal(inputData, &typeSpec); err != nil {
 			return err
 		}
 	}
@@ -1642,7 +1642,7 @@ func emptyObject(data RawMessage) (bool, error) {
 	return len(*objectData) == 0, nil
 }
 
-func UnmarshalFunctionSpec(funcSpec *FunctionSpec, data map[string]RawMessage) error {
+func unmarshalFunctionSpec(funcSpec *FunctionSpec, data map[string]RawMessage) error {
 	if description, ok := data["description"]; ok {
 		if err := json.Unmarshal(description, &funcSpec.Description); err != nil {
 			return err
@@ -1709,17 +1709,17 @@ func (funcSpec *FunctionSpec) UnmarshalJSON(inputJSON []byte) error {
 	if err := json.Unmarshal(inputJSON, &data); err != nil {
 		return err
 	}
-	return UnmarshalFunctionSpec(funcSpec, data)
+	return unmarshalFunctionSpec(funcSpec, data)
 }
 
 // UnmarshalYAML is custom unmarshalling logic for FunctionSpec so that we can derive Outputs from ReturnType
 // which otherwise isn't possible when both are retrieved from the same JSON field
-func (funcSpec *FunctionSpec) UnmarshalYAML(node *yaml.Node) error {
+func (funcSpec *FunctionSpec) UnmarshalYAML(input []byte) error {
 	var data map[string]RawMessage
-	if err := node.Decode(&data); err != nil {
+	if err := yaml.Unmarshal(input, &data); err != nil {
 		return err
 	}
-	return UnmarshalFunctionSpec(funcSpec, data)
+	return unmarshalFunctionSpec(funcSpec, data)
 }
 
 func (funcSpec FunctionSpec) MarshalFunctionSpec() map[string]interface{} {
diff --git a/pkg/codegen/schema/schema_test.go b/pkg/codegen/schema/schema_test.go
index 3aab8a3581ca..a0a326dfa13a 100644
--- a/pkg/codegen/schema/schema_test.go
+++ b/pkg/codegen/schema/schema_test.go
@@ -17,6 +17,7 @@ package schema
 
 import (
 	"encoding/json"
+	"fmt"
 	"io/ioutil"
 	"net/url"
 	"path/filepath"
@@ -27,6 +28,7 @@ import (
 
 	"github.com/blang/semver"
 	"github.com/stretchr/testify/assert"
+	"gopkg.in/yaml.v2"
 )
 
 func readSchemaFile(file string) (pkgSpec PackageSpec) {
@@ -36,8 +38,16 @@ func readSchemaFile(file string) (pkgSpec PackageSpec) {
 		panic(err)
 	}
 
-	if err = json.Unmarshal(schemaBytes, &pkgSpec); err != nil {
-		panic(err)
+	if strings.HasSuffix(file, ".json") {
+		if err = json.Unmarshal(schemaBytes, &pkgSpec); err != nil {
+			panic(err)
+		}
+	} else if strings.HasSuffix(file, ".yaml") || strings.HasSuffix(file, ".yml") {
+		if err = yaml.Unmarshal(schemaBytes, &pkgSpec); err != nil {
+			panic(err)
+		}
+	} else {
+		panic(fmt.Sprintf("unknown schema file extension while parsing %s", file))
 	}
 
 	return pkgSpec
@@ -402,6 +412,14 @@ func TestMethods(t *testing.T) {
 				assert.Equal(t, pkg.Functions[0].ReturnType, NumberType)
 			},
 		},
+		{
+			filename: "good-simplified-methods.yml",
+			validator: func(pkg *Package) {
+				assert.Len(t, pkg.Functions, 1)
+				assert.NotNil(t, pkg.Functions[0].ReturnType, "There should be a return type")
+				assert.Equal(t, pkg.Functions[0].ReturnType, NumberType)
+			},
+		},
 		{
 			filename:      "bad-methods-1.json",
 			expectedError: "unknown function xyz:index:Foo/bar",
diff --git a/pkg/codegen/testing/test/testdata/schema/good-simplified-methods.yml b/pkg/codegen/testing/test/testdata/schema/good-simplified-methods.yml
new file mode 100644
index 000000000000..31892c950390
--- /dev/null
+++ b/pkg/codegen/testing/test/testdata/schema/good-simplified-methods.yml
@@ -0,0 +1,14 @@
+name: xyz
+version: 0.0.1
+functions:
+  xyz:index:basic:
+    multiArgumentInputs:
+    - value
+    inputs:
+      properties:
+        value:
+          type: number
+      required:
+      - value
+    outputs:
+      type: number