feat: add UnmarshalYAML to protoutil (#61)

Author: yolocs

protoutil/conversion.go

@@ -19,10 +19,14 @@ import (
 	"encoding/json"
 	"fmt"
 
+	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/structpb"
+	"gopkg.in/yaml.v2"
 )
 
-// ToProtoStruct converts v, which must marshal into a JSON object, into a proto struct.
+// ToProtoStruct converts v, which must marshal into a JSON object, into a proto
+// struct.
 func ToProtoStruct(v any) (*structpb.Struct, error) {
 	jb, err := json.Marshal(v)
 	if err != nil {
@@ -34,3 +38,20 @@ func ToProtoStruct(v any) (*structpb.Struct, error) {
 	}
 	return x, nil
 }
+
+// UnmarshalYAML unmarshals the give YAML bytes to the given proto message.
+func UnmarshalYAML(b []byte, msg proto.Message) error {
+	tmp := map[string]any{}
+	if err := yaml.Unmarshal(b, tmp); err != nil {
+		return fmt.Errorf("failed to unmarshal yaml: %w", err)
+	}
+	jb, err := json.Marshal(tmp)
+	if err != nil {
+		return fmt.Errorf("failed to marshal json: %w", err)
+	}
+
+	if err := protojson.Unmarshal(jb, msg); err != nil {
+		return fmt.Errorf("failed to unmarshal proto: %w", err)
+	}
+	return nil
+}

protoutil/conversion_test.go

@@ -24,7 +24,7 @@ import (
 	"google.golang.org/protobuf/types/known/structpb"
 )
 
-func TestParseProject(t *testing.T) {
+func TestToProtoStruct(t *testing.T) {
 	t.Parallel()
 
 	cases := []struct {
@@ -82,3 +82,61 @@ func TestParseProject(t *testing.T) {
 		})
 	}
 }
+
+func TestUnmarshalYAML(t *testing.T) {
+	t.Parallel()
+
+	cases := []struct {
+		name          string
+		b             []byte
+		want          *structpb.Struct
+		wantErrSubstr string
+	}{
+		{
+			name: "success",
+			b: []byte(`foo: bar
+slice:
+- abc
+- xyz
+num: 1
+bool: true
+`),
+			want: &structpb.Struct{
+				Fields: map[string]*structpb.Value{
+					"foo": structpb.NewStringValue("bar"),
+					"slice": structpb.NewListValue(&structpb.ListValue{
+						Values: []*structpb.Value{
+							structpb.NewStringValue("abc"),
+							structpb.NewStringValue("xyz"),
+						},
+					}),
+					"num":  structpb.NewNumberValue(1),
+					"bool": structpb.NewBoolValue(true),
+				},
+			},
+		},
+		{
+			name:          "invalid_yaml_error",
+			b:             []byte("foobar: {}{}"),
+			want:          &structpb.Struct{},
+			wantErrSubstr: "failed to unmarshal yaml",
+		},
+	}
+
+	for _, tc := range cases {
+		tc := tc
+
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			var msg structpb.Struct
+			err := UnmarshalYAML(tc.b, &msg)
+			if diff := testutil.DiffErrString(err, tc.wantErrSubstr); diff != "" {
+				t.Errorf("unexpected error: %s", diff)
+			}
+			if diff := cmp.Diff(tc.want, &msg, protocmp.Transform()); diff != "" {
+				t.Errorf("UnmarshalYAML (-want,+got):\n%s", diff)
+			}
+		})
+	}
+}