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 6c6e995cb81085b2898d2a111392fd98c52cb333..386b40bd96839fb1eb045c20fb8b4b15c93a42f6 100644 (file)
--- 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 6603df2140a47447318a67af4b339dbe7d3ea923..a5a1cd78fabdcb8eda537ed152dda64d992cb055 100644 (file)
--- 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) {