save/restore request context in future

[folly.git] / folly / wangle / test / FutureTest.cpp

diff --git a/folly/wangle/test/FutureTest.cpp b/folly/wangle/test/FutureTest.cpp
index 8d1f95d3ce203fbda13641c188875f42ce4f656a..1ecca30ceda73e040316fd5ff62b27b2f8e7dad9 100644 (file)
--- a/folly/wangle/test/FutureTest.cpp
+++ b/folly/wangle/test/FutureTest.cpp
@@ -28,6 +28,9 @@
 #include <folly/wangle/Future.h>
 #include <folly/wangle/ManualExecutor.h>
 
+#include <folly/io/async/Request.h>
+
+using namespace folly;
 using namespace folly::wangle;
 using std::pair;
 using std::string;
@@ -908,3 +911,42 @@ TEST(Future, detachRace) {
   f.reset();
   t1.join();
 }
+
+class TestData : public RequestData {
+ public:
+  explicit TestData(int data) : data_(data) {}
+  virtual ~TestData() {}
+  int data_;
+};
+
+TEST(Future, context) {
+
+  // Start a new context
+  RequestContext::create();
+
+  EXPECT_EQ(nullptr, RequestContext::get()->getContextData("test"));
+
+  // Set some test data
+  RequestContext::get()->setContextData(
+    "test",
+    std::unique_ptr<TestData>(new TestData(10)));
+
+  // Start a future
+  Promise<void> p;
+  auto future = p.getFuture().then([&]{
+    // Check that the context followed the future
+    EXPECT_TRUE(RequestContext::get() != nullptr);
+    auto a = dynamic_cast<TestData*>(
+      RequestContext::get()->getContextData("test"));
+    auto data = a->data_;
+    EXPECT_EQ(10, data);
+  });
+
+  // Clear the context
+  RequestContext::setContext(nullptr);
+
+  EXPECT_EQ(nullptr, RequestContext::get()->getContextData("test"));
+
+  // Fulfil the promise
+  p.setValue();
+}