From: Andrii Grynenko <andrii@fb.com>
Date: Mon, 18 Jul 2016 20:26:59 +0000 (-0700)
Subject: Fix mocking to support multiple overrides
X-Git-Tag: 2016.07.26~35
X-Git-Url: http://plrg.eecs.uci.edu/git/?a=commitdiff_plain;h=46763a64f652c6218d253fb4307a9722587187e2;p=folly.git

Fix mocking to support multiple overrides

Summary: When installing the mock, we should make sure to remove singleton from the creation_order list.

Differential Revision: D3580725

fbshipit-source-id: dfb489de1be860ab639380644eab0b45a07a1450
---

diff --git a/folly/Singleton-inl.h b/folly/Singleton-inl.h
index 6c6e995c..386b40bd 100644
--- a/folly/Singleton-inl.h
+++ b/folly/Singleton-inl.h
@@ -71,6 +71,16 @@ void SingletonHolder<T>::registerSingletonMock(CreateFunc c, TeardownFunc t) {
   }
   destroyInstance();
 
+  {
+    RWSpinLock::WriteHolder wh(&vault_.mutex_);
+
+    auto it = std::find(
+        vault_.creation_order_.begin(), vault_.creation_order_.end(), type());
+    if (it != vault_.creation_order_.end()) {
+      vault_.creation_order_.erase(it);
+    }
+  }
+
   std::lock_guard<std::mutex> entry_lock(mutex_);
 
   create_ = std::move(c);
diff --git a/folly/test/SingletonTest.cpp b/folly/test/SingletonTest.cpp
index 6603df21..a5a1cd78 100644
--- a/folly/test/SingletonTest.cpp
+++ b/folly/test/SingletonTest.cpp
@@ -582,6 +582,18 @@ TEST(Singleton, MockTest) {
 
   // If serial_count value is the same, then singleton was not replaced.
   EXPECT_NE(serial_count_first, serial_count_mock);
+
+  // Override existing mock using make_mock one more time
+  SingletonMock<Watchdog>::make_mock();
+
+  EXPECT_EQ(vault.registeredSingletonCount(), 1);
+  int serial_count_mock2 = SingletonMock<Watchdog>::try_get()->serial_number;
+
+  // If serial_count value is the same, then singleton was not replaced.
+  EXPECT_NE(serial_count_first, serial_count_mock2);
+  EXPECT_NE(serial_count_mock, serial_count_mock2);
+
+  vault.destroyInstances();
 }
 
 TEST(Singleton, DoubleRegistrationLogging) {