fix(fieldbehavior): traverse messages and collections to clear nested…

… field behaviors TRA-1320
einride · Apr 15, 2024 · 3113ede · 3113ede
1 parent 1a8c18f
commit 3113ede
Show file tree

Hide file tree

Showing 4 changed files with 902 additions and 23 deletions.
diff --git a/fieldbehavior/fieldbehavior.go b/fieldbehavior/fieldbehavior.go
@@ -109,33 +109,72 @@ func isPresent(v protoreflect.Value, f protoreflect.FieldDescriptor) bool {
 }
 
 func clearFieldsWithBehaviors(m proto.Message, behaviorsToClear ...annotations.FieldBehavior) {
-	rangeFieldsWithBehaviors(m, func(
-		m protoreflect.Message,
-		f protoreflect.FieldDescriptor,
-		_ protoreflect.Value,
-		behaviors []annotations.FieldBehavior,
-	) bool {
-		if hasAnyBehavior(behaviors, behaviorsToClear) {
-			m.Clear(f)
-		}
-		return true
-	})
+	rangeFieldsWithBehaviors(
+		m.ProtoReflect(),
+		func(
+			m protoreflect.Message,
+			f protoreflect.FieldDescriptor,
+			_ protoreflect.Value,
+			behaviors []annotations.FieldBehavior,
+		) bool {
+			if hasAnyBehavior(behaviors, behaviorsToClear) {
+				m.Clear(f)
+			}
+			return true
+		},
+	)
 }
 
 func rangeFieldsWithBehaviors(
-	m proto.Message,
-	fn func(protoreflect.Message, protoreflect.FieldDescriptor, protoreflect.Value, []annotations.FieldBehavior) bool,
+	m protoreflect.Message,
+	fn func(
+		protoreflect.Message,
+		protoreflect.FieldDescriptor,
+		protoreflect.Value,
+		[]annotations.FieldBehavior,
+	) bool,
 ) {
-	d := m.ProtoReflect()
-	d.Range(func(f protoreflect.FieldDescriptor, v protoreflect.Value) bool {
-		if behaviors, ok := proto.GetExtension(
-			f.Options(),
-			annotations.E_FieldBehavior,
-		).([]annotations.FieldBehavior); ok {
-			return fn(d, f, v, behaviors)
-		}
-		return true
-	})
+	m.Range(
+		func(f protoreflect.FieldDescriptor, v protoreflect.Value) bool {
+			if behaviors, ok := proto.GetExtension(
+				f.Options(),
+				annotations.E_FieldBehavior,
+			).([]annotations.FieldBehavior); ok {
+				fn(m, f, v, behaviors)
+			}
+
+			switch {
+			// if field is repeated, traverse the nested message for field behaviors
+			case f.IsList() && f.Kind() == protoreflect.MessageKind:
+				for i := 0; i < v.List().Len(); i++ {
+					rangeFieldsWithBehaviors(
+						v.List().Get(i).Message(),
+						fn,
+					)
+				}
+				return true
+			// if field is map, traverse the nested message for field behaviors
+			case f.IsMap() && f.MapValue().Kind() == protoreflect.MessageKind:
+				v.Map().Range(func(_ protoreflect.MapKey, mv protoreflect.Value) bool {
+					rangeFieldsWithBehaviors(
+						mv.Message(),
+						fn,
+					)
+					return true
+				})
+				return true
+			// if field is message, traverse the message
+			// maps are also treated as Kind message and should not be traversed as messages
+			case f.Kind() == protoreflect.MessageKind && !f.IsMap():
+				rangeFieldsWithBehaviors(
+					v.Message(),
+					fn,
+				)
+				return true
+			default:
+				return true
+			}
+		})
 }
 
 func hasAnyBehavior(haystack, needles []annotations.FieldBehavior) bool {

diff --git a/fieldbehavior/fieldbehavior_test.go b/fieldbehavior/fieldbehavior_test.go
@@ -4,8 +4,10 @@ import (
 	"testing"
 
 	examplefreightv1 "go.einride.tech/aip/proto/gen/einride/example/freight/v1"
+	syntaxv1 "go.einride.tech/aip/proto/gen/einride/example/syntax/v1"
 	"google.golang.org/genproto/googleapis/api/annotations"
 	"google.golang.org/genproto/googleapis/example/library/v1"
+	"google.golang.org/protobuf/testing/protocmp"
 	"google.golang.org/protobuf/types/known/fieldmaskpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 	"gotest.tools/v3/assert"
@@ -26,6 +28,272 @@ func TestClearFields(t *testing.T) {
 		assert.Equal(t, site.GetDisplayName(), "site one")
 		assert.Equal(t, site.GetName(), "site1")
 	})
+	t.Run("clear field with set field_behavior on nested message", func(t *testing.T) {
+		t.Parallel()
+		input := &syntaxv1.FieldBehaviorMessage{
+			Field:         "field",    // has no field_behaviors; should not be cleared.
+			OptionalField: "optional", // has OPTIONAL field_behavior; should not be cleared.
+			MessageWithoutFieldBehavior: &syntaxv1.FieldBehaviorMessage{ // has no field_behaviors; should not be cleared.
+				Field:           "field",       // has no field_behaviors; should not be cleared.
+				OptionalField:   "optional",    // has OPTIONAL field_behavior; should not be cleared.
+				OutputOnlyField: "output_only", // has OUTPUT_ONLY field_behavior; should be cleared.
+			},
+		}
+
+		expected := &syntaxv1.FieldBehaviorMessage{
+			Field:         "field",
+			OptionalField: "optional",
+			MessageWithoutFieldBehavior: &syntaxv1.FieldBehaviorMessage{
+				Field:         "field",
+				OptionalField: "optional",
+			},
+		}
+
+		ClearFields(input, annotations.FieldBehavior_OUTPUT_ONLY)
+		assert.DeepEqual(t, input, expected, protocmp.Transform())
+	})
+
+	t.Run("clear field with set field_behavior on multiple levels of nested messages", func(t *testing.T) {
+		t.Parallel()
+		input := &syntaxv1.FieldBehaviorMessage{
+			Field:         "field",    // has no field_behaviors; should not be cleared.
+			OptionalField: "optional", // has OPTIONAL field_behavior; should not be cleared.
+			MessageWithoutFieldBehavior: &syntaxv1.FieldBehaviorMessage{ // has no field_behaviors; should not be cleared.
+				Field:           "field",       // has no field_behaviors; should not be cleared.
+				OptionalField:   "optional",    // has OPTIONAL field_behavior; should not be cleared.
+				OutputOnlyField: "output_only", // has OUTPUT_ONLY field_behavior; should be cleared.
+				MessageWithoutFieldBehavior: &syntaxv1.FieldBehaviorMessage{ // has no field_behaviors; should not be cleared.
+					Field:           "field",       // has no field_behaviors; should not be cleared.
+					OptionalField:   "optional",    // has OPTIONAL field_behavior; should not be cleared.
+					OutputOnlyField: "output_only", // has OUTPUT_ONLY field_behavior; should be cleared.
+				},
+			},
+		}
+
+		expected := &syntaxv1.FieldBehaviorMessage{
+			Field:         "field",
+			OptionalField: "optional",
+			MessageWithoutFieldBehavior: &syntaxv1.FieldBehaviorMessage{
+				Field:         "field",
+				OptionalField: "optional",
+				MessageWithoutFieldBehavior: &syntaxv1.FieldBehaviorMessage{
+					Field:         "field",
+					OptionalField: "optional",
+				},
+			},
+		}
+
+		ClearFields(input, annotations.FieldBehavior_OUTPUT_ONLY)
+		assert.DeepEqual(t, input, expected, protocmp.Transform())
+	})
+
+	t.Run("clear fields with set field_behavior on repeated message", func(t *testing.T) {
+		t.Parallel()
+		input := &syntaxv1.FieldBehaviorMessage{
+			Field:         "field",    // has no field_behaviors; should not be cleared.
+			OptionalField: "optional", // has OPTIONAL field_behavior; should not be cleared.
+			RepeatedMessage: []*syntaxv1.FieldBehaviorMessage{ // has no field_behaviors; should not be cleared.
+				{
+					Field:           "field",       // has no field_behaviors; should not be cleared.
+					OptionalField:   "optional",    // has OPTIONAL field_behavior; should not be cleared.
+					OutputOnlyField: "output_only", // has OUTPUT_ONLY field_behavior; should be cleared.
+				},
+			},
+			StringList: []string{ // not a message type, should not be traversed
+				"string",
+			},
+		}
+
+		expected := &syntaxv1.FieldBehaviorMessage{
+			Field:         "field",
+			OptionalField: "optional",
+			RepeatedMessage: []*syntaxv1.FieldBehaviorMessage{
+				{
+					Field:         "field",
+					OptionalField: "optional",
+				},
+			},
+			StringList: []string{
+				"string",
+			},
+		}
+
+		ClearFields(input, annotations.FieldBehavior_OUTPUT_ONLY)
+		assert.DeepEqual(t, input, expected, protocmp.Transform())
+	})
+
+	t.Run("clear fields with set field_behavior on multiple levels of repeated messages", func(t *testing.T) {
+		t.Parallel()
+		input := &syntaxv1.FieldBehaviorMessage{
+			Field:         "field",    // has no field_behaviors; should not be cleared.
+			OptionalField: "optional", // has OPTIONAL field_behavior; should not be cleared.
+			RepeatedMessage: []*syntaxv1.FieldBehaviorMessage{ // has no field_behaviors; should not be cleared.
+				{
+					Field:           "field",       // has no field_behaviors; should not be cleared.
+					OptionalField:   "optional",    // has OPTIONAL field_behavior; should not be cleared.
+					OutputOnlyField: "output_only", // has OUTPUT_ONLY field_behavior; should be cleared.
+					RepeatedMessage: []*syntaxv1.FieldBehaviorMessage{ // has no field_behaviors; should not be cleared.
+						{
+							Field:           "field",       // has no field_behaviors; should not be cleared.
+							OptionalField:   "optional",    // has OPTIONAL field_behavior; should not be cleared.
+							OutputOnlyField: "output_only", // has OUTPUT_ONLY field_behavior; should be cleared.
+						},
+					},
+				},
+			},
+			StringList: []string{ // not a message type, should not be traversed
+				"string",
+			},
+		}
+
+		expected := &syntaxv1.FieldBehaviorMessage{
+			Field:         "field",
+			OptionalField: "optional",
+			RepeatedMessage: []*syntaxv1.FieldBehaviorMessage{
+				{
+					Field:         "field",
+					OptionalField: "optional",
+					RepeatedMessage: []*syntaxv1.FieldBehaviorMessage{
+						{
+							Field:         "field",
+							OptionalField: "optional",
+						},
+					},
+				},
+			},
+			StringList: []string{
+				"string",
+			},
+		}
+
+		ClearFields(input, annotations.FieldBehavior_OUTPUT_ONLY)
+		assert.DeepEqual(t, input, expected, protocmp.Transform())
+	})
+
+	t.Run("clear repeated field with set field_behavior", func(t *testing.T) {
+		t.Parallel()
+		input := &syntaxv1.FieldBehaviorMessage{
+			Field:         "field",    // has no field_behaviors; should not be cleared.
+			OptionalField: "optional", // has OPTIONAL field_behavior; should not be cleared.
+			RepeatedMessage: []*syntaxv1.FieldBehaviorMessage{ // has no field_behaviors; should not be cleared.
+				{
+					Field:         "field",    // has no field_behaviors; should not be cleared.
+					OptionalField: "optional", // has OPTIONAL field_behavior; should not be cleared.
+				},
+			},
+			RepeatedOutputOnlyMessage: []*syntaxv1.FieldBehaviorMessage{ // has OUTPUT_ONLY field_behavior; should be cleared.
+				{
+					Field:           "field",
+					OptionalField:   "optional",
+					OutputOnlyField: "output_only",
+				},
+			},
+		}
+
+		expected := &syntaxv1.FieldBehaviorMessage{
+			Field:         "field",    // has no field_behaviors; should not be cleared.
+			OptionalField: "optional", // has OPTIONAL field_behavior; should not be cleared.
+			RepeatedMessage: []*syntaxv1.FieldBehaviorMessage{ // has no field_behaviors; should not be cleared.
+				{
+					Field:         "field",    // has no field_behaviors; should not be cleared.
+					OptionalField: "optional", // has OPTIONAL field_behavior; should not be cleared.
+				},
+			},
+		}
+
+		ClearFields(input, annotations.FieldBehavior_OUTPUT_ONLY)
+		assert.DeepEqual(t, input, expected, protocmp.Transform())
+	})
+
+	t.Run("clear fields with set field_behavior on message in map", func(t *testing.T) {
+		t.Parallel()
+		input := &syntaxv1.FieldBehaviorMessage{
+			OptionalField: "optional", // has OPTIONAL field_behavior; should not be cleared.
+			MapOptionalMessage: map[string]*syntaxv1.FieldBehaviorMessage{
+				"key_1": {
+					OptionalField:   "optional",    // has OPTIONAL field_behavior; should not be cleared.
+					OutputOnlyField: "output_only", // has OUTPUT_ONLY field_behavior; should be cleared.
+				},
+			},
+			StringMap: map[string]string{
+				"string_key": "string", // not a message type, should not be traversed
+			},
+		}
+
+		expected := &syntaxv1.FieldBehaviorMessage{
+			OptionalField: "optional",
+			MapOptionalMessage: map[string]*syntaxv1.FieldBehaviorMessage{
+				"key_1": {
+					OptionalField: "optional",
+				},
+			},
+			StringMap: map[string]string{
+				"string_key": "string",
+			},
+		}
+
+		ClearFields(input, annotations.FieldBehavior_OUTPUT_ONLY)
+		assert.DeepEqual(
+			t,
+			input,
+			expected,
+			protocmp.Transform(),
+		)
+	})
+
+	t.Run("clear map field with set field_behavior", func(t *testing.T) {
+		t.Parallel()
+		input := &syntaxv1.FieldBehaviorMessage{
+			OptionalField: "optional",
+			// has OUTPUT_ONLY field_behavior; should be cleared.
+			MapOutputOnlyMessage: map[string]*syntaxv1.FieldBehaviorMessage{
+				"key_1": {
+					OutputOnlyField: "output_only", // has OUTPUT_ONLY field_behavior; but should be cleared with parent message.
+				},
+			},
+		}
+
+		expected := &syntaxv1.FieldBehaviorMessage{
+			OptionalField: "optional",
+		}
+
+		ClearFields(input, annotations.FieldBehavior_OUTPUT_ONLY)
+		assert.DeepEqual(
+			t,
+			input,
+			expected,
+			protocmp.Transform(),
+		)
+	})
+
+	t.Run("clear field with set field_behavior on oneof message", func(t *testing.T) {
+		t.Parallel()
+		input := &syntaxv1.FieldBehaviorMessage{
+			Field:         "field",    // has no field_behaviors; should not be cleared.
+			OptionalField: "optional", // has OPTIONAL field_behavior; should not be cleared.
+			Oneof: &syntaxv1.FieldBehaviorMessage_FieldBehaviorMessage{
+				FieldBehaviorMessage: &syntaxv1.FieldBehaviorMessage{
+					Field:           "field",       // has no field_behaviors; should not be cleared.
+					OptionalField:   "optional",    // has OPTIONAL field_behavior; should not be cleared.
+					OutputOnlyField: "output_only", // has OUTPUT_ONLY field_behavior; should be cleared.
+				},
+			},
+		}
+
+		expected := &syntaxv1.FieldBehaviorMessage{
+			Field:         "field",
+			OptionalField: "optional",
+			Oneof: &syntaxv1.FieldBehaviorMessage_FieldBehaviorMessage{
+				FieldBehaviorMessage: &syntaxv1.FieldBehaviorMessage{
+					Field:         "field",
+					OptionalField: "optional",
+				},
+			},
+		}
+
+		ClearFields(input, annotations.FieldBehavior_OUTPUT_ONLY)
+		assert.DeepEqual(t, input, expected, protocmp.Transform())
+	})
 }
 
 func TestCopyFields(t *testing.T) {