1 files changed, 47 insertions, 0 deletions
diff --git a/python/google/protobuf/internal/descriptor_pool_test.py b/python/google/protobuf/internal/descriptor_pool_test.py
index d4de2d81..cb6abe6c 100644
--- a/python/google/protobuf/internal/descriptor_pool_test.py
+++ b/python/google/protobuf/internal/descriptor_pool_test.py
@@ -254,6 +254,53 @@ class DescriptorPoolTest(unittest.TestCase):
     with self.assertRaises(KeyError):
       self.pool.FindFieldByName('Does not exist')
 
+  def testFindAllExtensions(self):
+    factory1_message = self.pool.FindMessageTypeByName(
+        'google.protobuf.python.internal.Factory1Message')
+    factory2_message = self.pool.FindMessageTypeByName(
+        'google.protobuf.python.internal.Factory2Message')
+    # An extension defined in a message.
+    one_more_field = factory2_message.extensions_by_name['one_more_field']
+    self.pool.AddExtensionDescriptor(one_more_field)
+    # An extension defined at file scope.
+    factory_test2 = self.pool.FindFileByName(
+        'google/protobuf/internal/factory_test2.proto')
+    another_field = factory_test2.extensions_by_name['another_field']
+    self.pool.AddExtensionDescriptor(another_field)
+
+    extensions = self.pool.FindAllExtensions(factory1_message)
+    expected_extension_numbers = {one_more_field, another_field}
+    self.assertEqual(expected_extension_numbers, set(extensions))
+    # Verify that mutating the returned list does not affect the pool.
+    extensions.append('unexpected_element')
+    # Get the extensions again, the returned value does not contain the
+    # 'unexpected_element'.
+    extensions = self.pool.FindAllExtensions(factory1_message)
+    self.assertEqual(expected_extension_numbers, set(extensions))
+
+  def testFindExtensionByNumber(self):
+    factory1_message = self.pool.FindMessageTypeByName(
+        'google.protobuf.python.internal.Factory1Message')
+    factory2_message = self.pool.FindMessageTypeByName(
+        'google.protobuf.python.internal.Factory2Message')
+    # An extension defined in a message.
+    one_more_field = factory2_message.extensions_by_name['one_more_field']
+    self.pool.AddExtensionDescriptor(one_more_field)
+    # An extension defined at file scope.
+    factory_test2 = self.pool.FindFileByName(
+        'google/protobuf/internal/factory_test2.proto')
+    another_field = factory_test2.extensions_by_name['another_field']
+    self.pool.AddExtensionDescriptor(another_field)
+
+    # An extension defined in a message.
+    extension = self.pool.FindExtensionByNumber(factory1_message, 1001)
+    self.assertEqual(extension.name, 'one_more_field')
+    # An extension defined at file scope.
+    extension = self.pool.FindExtensionByNumber(factory1_message, 1002)
+    self.assertEqual(extension.name, 'another_field')
+    with self.assertRaises(KeyError):
+      extension = self.pool.FindExtensionByNumber(factory1_message, 1234567)
+
   def testExtensionsAreNotFields(self):
     with self.assertRaises(KeyError):
       self.pool.FindFieldByName('google.protobuf.python.internal.another_field')