testing/protocmp: fix reflection handling of extensions

Extensions should be checked based on ContainingMessage,
rather than the Parent. Add tests to ensure this works.

Change-Id: Iaf257f65197fb8d332039bc77a192753f8c4159f
Reviewed-on: https://go-review.googlesource.com/c/protobuf/+/221426
Reviewed-by: Damien Neil <dneil@google.com>
diff --git a/testing/protocmp/reflect.go b/testing/protocmp/reflect.go
index d1677f6..62c646f 100644
--- a/testing/protocmp/reflect.go
+++ b/testing/protocmp/reflect.go
@@ -38,7 +38,7 @@
 type reflectMessage Message
 
 func (m reflectMessage) stringKey(fd protoreflect.FieldDescriptor) string {
-	if m.Descriptor() != fd.Parent() {
+	if m.Descriptor() != fd.ContainingMessage() {
 		panic("mismatching containing message")
 	}
 	if fd.IsExtension() {
diff --git a/testing/protocmp/reflect_test.go b/testing/protocmp/reflect_test.go
index e9f3898..07489e9 100644
--- a/testing/protocmp/reflect_test.go
+++ b/testing/protocmp/reflect_test.go
@@ -89,6 +89,36 @@
 		&testpb.TestAllTypes{
 			OneofField: &testpb.TestAllTypes_OneofEnum{testpb.TestAllTypes_NEG},
 		},
+		func() proto.Message {
+			m := new(testpb.TestAllExtensions)
+			proto.SetExtension(m, testpb.E_OptionalInt32, int32(-32))
+			proto.SetExtension(m, testpb.E_OptionalInt64, int64(-64))
+			proto.SetExtension(m, testpb.E_OptionalUint32, uint32(32))
+			proto.SetExtension(m, testpb.E_OptionalUint64, uint64(64))
+			proto.SetExtension(m, testpb.E_OptionalFloat, float32(32.32))
+			proto.SetExtension(m, testpb.E_OptionalDouble, float64(64.64))
+			proto.SetExtension(m, testpb.E_OptionalBool, bool(true))
+			proto.SetExtension(m, testpb.E_OptionalString, string("string"))
+			proto.SetExtension(m, testpb.E_OptionalBytes, []byte("bytes"))
+			proto.SetExtension(m, testpb.E_OptionalNestedMessage, &testpb.TestAllExtensions_NestedMessage{A: proto.Int32(-32)})
+			proto.SetExtension(m, testpb.E_OptionalNestedEnum, testpb.TestAllTypes_NEG)
+			return m
+		}(),
+		func() proto.Message {
+			m := new(testpb.TestAllExtensions)
+			proto.SetExtension(m, testpb.E_RepeatedInt32, []int32{-32, +32})
+			proto.SetExtension(m, testpb.E_RepeatedInt64, []int64{-64, +64})
+			proto.SetExtension(m, testpb.E_RepeatedUint32, []uint32{0, 32})
+			proto.SetExtension(m, testpb.E_RepeatedUint64, []uint64{0, 64})
+			proto.SetExtension(m, testpb.E_RepeatedFloat, []float32{-32.32, +32.32})
+			proto.SetExtension(m, testpb.E_RepeatedDouble, []float64{-64.64, +64.64})
+			proto.SetExtension(m, testpb.E_RepeatedBool, []bool{false, true})
+			proto.SetExtension(m, testpb.E_RepeatedString, []string{"hello", "goodbye"})
+			proto.SetExtension(m, testpb.E_RepeatedBytes, [][]byte{[]byte("hello"), []byte("goodbye")})
+			proto.SetExtension(m, testpb.E_RepeatedNestedMessage, []*testpb.TestAllExtensions_NestedMessage{{A: proto.Int32(-32)}, {A: proto.Int32(+32)}})
+			proto.SetExtension(m, testpb.E_RepeatedNestedEnum, []testpb.TestAllTypes_NestedEnum{testpb.TestAllTypes_FOO, testpb.TestAllTypes_NEG})
+			return m
+		}(),
 		&textpb.KnownTypes{
 			OptBool:   &wrapperspb.BoolValue{Value: true},
 			OptInt32:  &wrapperspb.Int32Value{Value: -32},
diff --git a/testing/protocmp/xform.go b/testing/protocmp/xform.go
index a42bf9b..3b2f9f3 100644
--- a/testing/protocmp/xform.go
+++ b/testing/protocmp/xform.go
@@ -127,10 +127,6 @@
 	panic("invalid mutation of a read-only message")
 }
 
-// TODO: There is currently no public API for retrieving the FieldDescriptors
-// for extension fields. Rather than adding a specialized API to support that,
-// perhaps Message should just implement protoreflect.ProtoMessage instead.
-
 // String returns a formatted string for the message.
 // It is intended for human debugging and has no guarantees about its
 // exact format or the stability of its output.