Tests for marshalling JSON/YAML for FunctionSpec and refactor codegen

pkg/codegen/dotnet/gen.go

@@ -1383,26 +1383,23 @@ func (mod *modContext) functionReturnType(fun *schema.Function) string {
 //
 // We use Invoke for functions with object return types and InvokeSingle for everything else.
 func runtimeInvokeFunction(fun *schema.Function) string {
+	switch fun.ReturnType.(type) {
 	// If the function has no return type, it is a void function.
-	if fun.ReturnType == nil {
+	case nil:
 		return "Invoke"
-	}
-
 	// If the function has an object return type, it is a normal invoke function.
-	if _, isObject := fun.ReturnType.(*schema.ObjectType); isObject {
+	case *schema.ObjectType:
 		return "Invoke"
-	}
-
 	// If the function has an object return type, it is also a normal invoke function.
 	// because the deserialization can handle it
-	if _, isMap := fun.ReturnType.(*schema.MapType); isMap {
+	case *schema.MapType:
 		return "Invoke"
+	default:
+		// Anything else needs to be handled by InvokeSingle
+		// which expects an object with a single property to be returned
+		// then unwraps the value from that property
+		return "InvokeSingle"
 	}
-
-	// Anything else needs to be handled by InvokeSingle
-	// which expects an object with a single property to be returned
-	// then unwraps the value from that property
-	return "InvokeSingle"
 }
 
 func (mod *modContext) genFunction(w io.Writer, fun *schema.Function) error {

pkg/codegen/nodejs/gen.go

@@ -1109,31 +1109,28 @@ func (mod *modContext) functionReturnType(fun *schema.Function) string {
 // runtimeInvokeFunction returns the name of the Invoke function to use at runtime
 // from the SDK for the given provider function. This is necessary because some
 // functions have simple return types such as number, string, array<string> etc.
-// and the SDK's Invoke function cannot handle these types since the engine expects
+// and the SDK's invoke function cannot handle these types since the engine expects
 // the result of invokes to be a dictionary.
 //
-// We use Invoke for functions with object return types and InvokeSingle for everything else.
+// We use invoke for functions with object return types and invokeSingle for everything else.
 func runtimeInvokeFunction(fun *schema.Function) string {
+	switch fun.ReturnType.(type) {
 	// If the function has no return type, it is a void function.
-	if fun.ReturnType == nil {
+	case nil:
 		return "invoke"
-	}
-
 	// If the function has an object return type, it is a normal invoke function.
-	if _, isObject := fun.ReturnType.(*schema.ObjectType); isObject {
+	case *schema.ObjectType:
 		return "invoke"
-	}
-
 	// If the function has an object return type, it is also a normal invoke function.
 	// because the deserialization can handle it
-	if _, isMap := fun.ReturnType.(*schema.MapType); isMap {
+	case *schema.MapType:
 		return "invoke"
+	default:
+		// Anything else needs to be handled by InvokeSingle
+		// which expects an object with a single property to be returned
+		// then unwraps the value from that property
+		return "invokeSingle"
 	}
-
-	// Anything else needs to be handled by InvokeSingle
-	// which expects an object with a single property to be returned
-	// then unwraps the value from that property
-	return "invokeSingle"
 }
 
 func (mod *modContext) genFunction(w io.Writer, fun *schema.Function) (functionFileInfo, error) {
@@ -1168,14 +1165,10 @@ func (mod *modContext) genFunction(w io.Writer, fun *schema.Function) (functionF
 				fmt.Fprintf(w, "%s, ", mod.typeString(prop.Type, false, nil))
 			} else {
 				fmt.Fprintf(w, "%s?: ", prop.Name)
-				// since we already applied the ? to the type, we can simplify
+				// since we already applied the '?' to the type, we can simplify
 				// the optional-ness of the type
-				switch propType := prop.Type.(type) {
-				case *schema.OptionalType:
-					fmt.Fprintf(w, "%s, ", mod.typeString(propType.ElementType, false, nil))
-				default:
-					fmt.Fprintf(w, "%s, ", mod.typeString(propType, false, nil))
-				}
+				propType := prop.Type.(*schema.OptionalType)
+				fmt.Fprintf(w, "%s, ", mod.typeString(propType.ElementType, false, nil))
 			}
 		}
 	} else {
@@ -1203,7 +1196,7 @@ func (mod *modContext) genFunction(w io.Writer, fun *schema.Function) (functionF
 			// Pass the argument to the invocation.
 			body := fmt.Sprintf("args.%s", p.Name)
 			if fun.MultiArgumentInputs {
-				body = fmt.Sprintf("%s", p.Name)
+				body = p.Name
 			}
 
 			if name := mod.provideDefaultsFuncName(p.Type, true /*input*/); name != "" {

pkg/codegen/schema/schema.go

@@ -1552,24 +1552,28 @@ type ReturnTypeSpec struct {
 	TypeSpec       *TypeSpec
 }
 
-func (returnTypeSpec *ReturnTypeSpec) UnmarshalJSON(inputJSON []byte) error {
-	var data map[string]interface{}
-	if err := json.Unmarshal(inputJSON, &data); err != nil {
+// Decoder is an alias for a function that takes (in []byte, out interface{}) and potentially returns an error
+// it is used to abstract json.Unmarshal and yaml.Unmarshal which satisfy this function signature
+type Decoder func([]byte, interface{}) error
+
+func (returnTypeSpec *ReturnTypeSpec) UnmarshalReturnTypeSpec(data []byte, decode Decoder) error {
+	var objectMap map[string]interface{}
+	if err := decode(data, &objectMap); err != nil {
 		return err
 	}
 
-	if len(data) == 0 {
+	if len(objectMap) == 0 {
 		return nil
 	}
 
 	var objectSpec *ObjectTypeSpec
 	var typeSpec *TypeSpec
-	if _, hasProperties := data["properties"]; hasProperties {
-		if err := json.Unmarshal(inputJSON, &objectSpec); err != nil {
+	if _, hasProperties := objectMap["properties"]; hasProperties {
+		if err := decode(data, &objectSpec); err != nil {
 			return err
 		}
 	} else {
-		if err := json.Unmarshal(inputJSON, &typeSpec); err != nil {
+		if err := decode(data, &typeSpec); err != nil {
 			return err
 		}
 	}
@@ -1579,31 +1583,12 @@ func (returnTypeSpec *ReturnTypeSpec) UnmarshalJSON(inputJSON []byte) error {
 	return nil
 }
 
-func (returnTypeSpec *ReturnTypeSpec) UnmarshalYAML(inputData []byte) error {
-	var data map[string]interface{}
-	if err := yaml.Unmarshal(inputData, &data); err != nil {
-		return err
-	}
-
-	if len(data) == 0 {
-		return nil
-	}
-
-	var objectSpec *ObjectTypeSpec
-	var typeSpec *TypeSpec
-	if _, hasProperties := data["properties"]; hasProperties {
-		if err := yaml.Unmarshal(inputData, &objectSpec); err != nil {
-			return err
-		}
-	} else {
-		if err := yaml.Unmarshal(inputData, &typeSpec); err != nil {
-			return err
-		}
-	}
+func (returnTypeSpec *ReturnTypeSpec) UnmarshalJSON(inputJSON []byte) error {
+	return returnTypeSpec.UnmarshalReturnTypeSpec(inputJSON, json.Unmarshal)
+}
 
-	returnTypeSpec.TypeSpec = typeSpec
-	returnTypeSpec.ObjectTypeSpec = objectSpec
-	return nil
+func (returnTypeSpec *ReturnTypeSpec) UnmarshalYAML(inputYAML []byte) error {
+	return returnTypeSpec.UnmarshalReturnTypeSpec(inputYAML, yaml.Unmarshal)
 }
 
 // FunctionSpec is the serializable form of a function description.
@@ -1714,15 +1699,15 @@ func (funcSpec *FunctionSpec) UnmarshalJSON(inputJSON []byte) error {
 
 // 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(input []byte) error {
+func (funcSpec *FunctionSpec) UnmarshalYAML(node *yaml.Node) error {
 	var data map[string]RawMessage
-	if err := yaml.Unmarshal(input, &data); err != nil {
+	if err := node.Decode(&data); err != nil {
 		return err
 	}
 	return unmarshalFunctionSpec(funcSpec, data)
 }
 
-func (funcSpec FunctionSpec) MarshalFunctionSpec() map[string]interface{} {
+func (funcSpec FunctionSpec) marshalFunctionSpec() map[string]interface{} {
 	data := make(map[string]interface{})
 	if funcSpec.Description != "" {
 		data["description"] = funcSpec.Description
@@ -1768,11 +1753,11 @@ func (funcSpec FunctionSpec) MarshalFunctionSpec() map[string]interface{} {
 }
 
 func (funcSpec FunctionSpec) MarshalJSON() ([]byte, error) {
-	return json.Marshal(funcSpec.MarshalFunctionSpec())
+	return json.Marshal(funcSpec.marshalFunctionSpec())
 }
 
-func (funcSpec FunctionSpec) MarshalYAML() ([]byte, error) {
-	return yaml.Marshal(funcSpec.MarshalFunctionSpec())
+func (funcSpec FunctionSpec) MarshalYAML() (interface{}, error) {
+	return funcSpec.marshalFunctionSpec(), nil
 }
 
 // ConfigSpec is the serializable description of a package's configuration variables.

pkg/codegen/schema/schema_test.go

@@ -28,7 +28,7 @@ import (
 
 	"github.com/blang/semver"
 	"github.com/stretchr/testify/assert"
-	"gopkg.in/yaml.v2"
+	"gopkg.in/yaml.v3"
 )
 
 func readSchemaFile(file string) (pkgSpec PackageSpec) {
@@ -119,6 +119,70 @@ var enumTests = []struct {
 	}},
 }
 
+func TestUnmarshalYAMLFunctionSpec(t *testing.T) {
+	t.Parallel()
+	var functionSpec *FunctionSpec
+	fnYaml := `
+description: Test function
+outputs:
+  type: number`
+
+	err := yaml.Unmarshal([]byte(fnYaml), &functionSpec)
+	assert.Nil(t, err, "Unmarshalling should work")
+	assert.Equal(t, "Test function", functionSpec.Description)
+	assert.NotNil(t, functionSpec.ReturnType, "Return type is not nil")
+	assert.NotNil(t, functionSpec.ReturnType.TypeSpec, "Return type is a type spec")
+	assert.Equal(t, "number", functionSpec.ReturnType.TypeSpec.Type, "Return type is a number")
+}
+
+func TestUnmarshalJSONFunctionSpec(t *testing.T) {
+	t.Parallel()
+	var functionSpec *FunctionSpec
+	fnJSON := `{"description":"Test function", "outputs": { "type": "number" } }`
+	err := json.Unmarshal([]byte(fnJSON), &functionSpec)
+	assert.Nil(t, err, "Unmarshalling should work")
+	assert.Equal(t, "Test function", functionSpec.Description)
+	assert.NotNil(t, functionSpec.ReturnType, "Return type is not nil")
+	assert.NotNil(t, functionSpec.ReturnType.TypeSpec, "Return type is a type spec")
+	assert.Equal(t, "number", functionSpec.ReturnType.TypeSpec.Type, "Return type is a number")
+}
+
+func TestMarshalJSONFunctionSpec(t *testing.T) {
+	t.Parallel()
+	functionSpec := &FunctionSpec{
+		Description: "Test function",
+		ReturnType: &ReturnTypeSpec{
+			TypeSpec: &TypeSpec{Type: "number"},
+		},
+	}
+
+	dataJSON, err := json.Marshal(functionSpec)
+	data := string(dataJSON)
+	expectedJSON := `{"description":"Test function","outputs":{"type":"number"}}`
+	assert.Nil(t, err, "Unmarshalling should work")
+	assert.Equal(t, expectedJSON, data)
+}
+
+func TestMarshalYAMLFunctionSpec(t *testing.T) {
+	t.Parallel()
+	functionSpec := &FunctionSpec{
+		Description: "Test function",
+		ReturnType: &ReturnTypeSpec{
+			TypeSpec: &TypeSpec{Type: "number"},
+		},
+	}
+
+	dataYAML, err := yaml.Marshal(functionSpec)
+	data := string(dataYAML)
+	expectedYAML := `description: Test function
+outputs:
+    type: number
+`
+
+	assert.Nil(t, err, "Unmarshalling should work")
+	assert.Equal(t, expectedYAML, data)
+}
+
 func TestEnums(t *testing.T) {
 	t.Parallel()