From d64a2d9941c36a7bc2a7959ea10ab8363192ac14 Mon Sep 17 00:00:00 2001
From: Adam Cozzette <acozzette@google.com>
Date: Wed, 29 Jun 2016 15:23:27 -0700
Subject: Integrated internal changes from Google

This includes all internal changes from around May 20 to now.
---
 src/google/protobuf/extension_set_unittest.cc | 68 +++++++++++++++++++++++++++
 1 file changed, 68 insertions(+)

(limited to 'src/google/protobuf/extension_set_unittest.cc')

diff --git a/src/google/protobuf/extension_set_unittest.cc b/src/google/protobuf/extension_set_unittest.cc
index f40fcbc2..688afedb 100644
--- a/src/google/protobuf/extension_set_unittest.cc
+++ b/src/google/protobuf/extension_set_unittest.cc
@@ -205,6 +205,74 @@ TEST(ExtensionSetTest, ReleaseExtension) {
   delete released_extension;
 }
 
+TEST(ExtensionSetTest, ArenaUnsafeArenaSetAllocatedAndRelease) {
+  ::google::protobuf::Arena arena;
+  unittest::TestAllExtensions* message =
+      ::google::protobuf::Arena::CreateMessage<unittest::TestAllExtensions>(&arena);
+  unittest::ForeignMessage extension;
+  message->UnsafeArenaSetAllocatedExtension(
+      unittest::optional_foreign_message_extension,
+      &extension);
+  // No copy when set.
+  unittest::ForeignMessage* mutable_extension =
+      message->MutableExtension(unittest::optional_foreign_message_extension);
+  EXPECT_EQ(&extension, mutable_extension);
+  // No copy when unsafe released.
+  unittest::ForeignMessage* released_extension =
+      message->UnsafeArenaReleaseExtension(
+          unittest::optional_foreign_message_extension);
+  EXPECT_EQ(&extension, released_extension);
+  EXPECT_FALSE(message->HasExtension(
+      unittest::optional_foreign_message_extension));
+  // Set the ownership back and let the destructors run.  It should not take
+  // ownership, so this should not crash.
+  message->UnsafeArenaSetAllocatedExtension(
+      unittest::optional_foreign_message_extension,
+      &extension);
+}
+
+TEST(ExtensionSetTest, UnsafeArenaSetAllocatedAndRelease) {
+  unittest::TestAllExtensions message;
+  unittest::ForeignMessage* extension = new unittest::ForeignMessage();
+  message.UnsafeArenaSetAllocatedExtension(
+      unittest::optional_foreign_message_extension,
+      extension);
+  // No copy when set.
+  unittest::ForeignMessage* mutable_extension =
+      message.MutableExtension(unittest::optional_foreign_message_extension);
+  EXPECT_EQ(extension, mutable_extension);
+  // No copy when unsafe released.
+  unittest::ForeignMessage* released_extension =
+      message.UnsafeArenaReleaseExtension(
+          unittest::optional_foreign_message_extension);
+  EXPECT_EQ(extension, released_extension);
+  EXPECT_FALSE(message.HasExtension(
+      unittest::optional_foreign_message_extension));
+  // Set the ownership back and let the destructors run.  It should take
+  // ownership, so this should not leak.
+  message.UnsafeArenaSetAllocatedExtension(
+      unittest::optional_foreign_message_extension,
+      extension);
+}
+
+TEST(ExtensionSetTest, ArenaUnsafeArenaReleaseOfHeapAlloc) {
+  ::google::protobuf::Arena arena;
+  unittest::TestAllExtensions* message =
+      ::google::protobuf::Arena::CreateMessage<unittest::TestAllExtensions>(&arena);
+  unittest::ForeignMessage* extension = new unittest::ForeignMessage;
+  message->SetAllocatedExtension(
+      unittest::optional_foreign_message_extension,
+      extension);
+  // The arena should maintain ownership of the heap allocated proto because we
+  // used UnsafeArenaReleaseExtension.  The leak checker will ensure this.
+  unittest::ForeignMessage* released_extension =
+      message->UnsafeArenaReleaseExtension(
+          unittest::optional_foreign_message_extension);
+  EXPECT_EQ(extension, released_extension);
+  EXPECT_FALSE(message->HasExtension(
+      unittest::optional_foreign_message_extension));
+}
+
 
 TEST(ExtensionSetTest, CopyFrom) {
   unittest::TestAllExtensions message1, message2;
-- 
cgit v1.2.3