From a2842c231e12c4c1913134c427d16084469905de Mon Sep 17 00:00:00 2001
From: heng <heng.chen@vesoft.com>
Date: Thu, 18 Jul 2019 15:30:53 +0800
Subject: [PATCH] Remove sleep

---
 src/kvstore/raftex/test/RaftexTestBase.cpp |  4 +---
 src/meta/test/TestUtils.h                  | 19 ++++++++++++-------
 src/storage/test/QueryBoundTest.cpp        |  4 +++-
 src/storage/test/QueryEdgePropsTest.cpp    |  4 +++-
 src/storage/test/QueryStatsTest.cpp        |  4 +++-
 src/storage/test/QueryVertexPropsTest.cpp  |  5 +++--
 src/storage/test/TestUtils.h               |  1 +
 7 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/src/kvstore/raftex/test/RaftexTestBase.cpp b/src/kvstore/raftex/test/RaftexTestBase.cpp
index 21c50240638..8cb8b48e6b2 100644
--- a/src/kvstore/raftex/test/RaftexTestBase.cpp
+++ b/src/kvstore/raftex/test/RaftexTestBase.cpp
@@ -100,7 +100,7 @@ void waitUntilLeaderElected(
             });
 
             // Sleep some time to wait until resp of heartbeat has come back when elected as leader
-            usleep(10000);
+            usleep(30000);
 
             bool sameLeader = true;
             int32_t index = 0;
@@ -121,8 +121,6 @@ void waitUntilLeaderElected(
         // Wait for one second
         sleep(1);
     }
-    // Wait for 50ms for the hb
-    usleep(50000);
 }
 
 
diff --git a/src/meta/test/TestUtils.h b/src/meta/test/TestUtils.h
index 5873919b645..b2e7d2e4b87 100644
--- a/src/meta/test/TestUtils.h
+++ b/src/meta/test/TestUtils.h
@@ -16,6 +16,7 @@
 #include "meta/processors/partsMan/ListHostsProcessor.h"
 #include "meta/MetaServiceHandler.h"
 #include <thrift/lib/cpp2/server/ThriftServer.h>
+#include <folly/synchronization/Baton.h>
 #include "meta/processors/usersMan/AuthenticationProcessor.h"
 #include "interface/gen-cpp2/common_types.h"
 #include "time/WallClock.h"
@@ -117,12 +118,13 @@ class TestUtils {
             data.emplace_back(MetaServiceUtils::partKey(id, partId),
                               MetaServiceUtils::partVal(hosts));
         }
-
+        folly::Baton<true, std::atomic> baton;
         kv->asyncMultiPut(0, 0, std::move(data),
                           [&] (kvstore::ResultCode code) {
                               ret = (code == kvstore::ResultCode::SUCCEEDED);
+                              baton.post();
                           });
-        usleep(10000);
+        baton.wait();
         return ret;
     }
 
@@ -144,12 +146,13 @@ class TestUtils {
             tags.emplace_back(MetaServiceUtils::schemaTagKey(1, tagId, ver++),
                               MetaServiceUtils::schemaTagVal(tagName, srcsch));
         }
-
+        folly::Baton<true, std::atomic> baton;
         kv->asyncMultiPut(0, 0, std::move(tags),
-                          [] (kvstore::ResultCode code) {
+                          [&] (kvstore::ResultCode code) {
                                 ASSERT_EQ(kvstore::ResultCode::SUCCEEDED, code);
+                                baton.post();
                           });
-        usleep(10000);
+        baton.wait();
     }
 
     static void mockEdge(kvstore::KVStore* kv, int32_t edgeNum, SchemaVer version = 0) {
@@ -172,10 +175,12 @@ class TestUtils {
                                MetaServiceUtils::schemaEdgeVal(edgeName, srcsch));
         }
 
-        kv->asyncMultiPut(0, 0, std::move(edges), [] (kvstore::ResultCode code) {
+        folly::Baton<true, std::atomic> baton;
+        kv->asyncMultiPut(0, 0, std::move(edges), [&] (kvstore::ResultCode code) {
             ASSERT_EQ(kvstore::ResultCode::SUCCEEDED, code);
+            baton.post();
         });
-        usleep(10000);
+        baton.wait();
     }
 
     static std::unique_ptr<test::ServerContext> mockMetaServer(uint16_t port,
diff --git a/src/storage/test/QueryBoundTest.cpp b/src/storage/test/QueryBoundTest.cpp
index 7a991c7425b..d1cbc266df2 100644
--- a/src/storage/test/QueryBoundTest.cpp
+++ b/src/storage/test/QueryBoundTest.cpp
@@ -67,13 +67,15 @@ void mockData(kvstore::KVStore* kv) {
                 }
             }
         }
+        folly::Baton<true, std::atomic> baton;
         kv->asyncMultiPut(
             0, partId, std::move(data),
             [&](kvstore::ResultCode code) {
                 EXPECT_EQ(code, kvstore::ResultCode::SUCCEEDED);
+                baton.post();
             });
+        baton.wait();
     }
-    sleep(1);
 }
 
 
diff --git a/src/storage/test/QueryEdgePropsTest.cpp b/src/storage/test/QueryEdgePropsTest.cpp
index 7409f0e7db6..9662681cb54 100644
--- a/src/storage/test/QueryEdgePropsTest.cpp
+++ b/src/storage/test/QueryEdgePropsTest.cpp
@@ -36,13 +36,15 @@ void mockData(kvstore::KVStore* kv) {
                 data.emplace_back(std::move(key), std::move(val));
             }
         }
+        folly::Baton<true, std::atomic> baton;
         kv->asyncMultiPut(
             0, partId, std::move(data),
             [&](kvstore::ResultCode code) {
                 EXPECT_EQ(code, kvstore::ResultCode::SUCCEEDED);
+                baton.post();
             });
+        baton.wait();
     }
-    sleep(1);
 }
 
 
diff --git a/src/storage/test/QueryStatsTest.cpp b/src/storage/test/QueryStatsTest.cpp
index 68bdb4e88b9..cb87829fe0e 100644
--- a/src/storage/test/QueryStatsTest.cpp
+++ b/src/storage/test/QueryStatsTest.cpp
@@ -48,11 +48,13 @@ void mockData(kvstore::KVStore* kv) {
                 data.emplace_back(std::move(key), std::move(val));
             }
         }
+        folly::Baton<true, std::atomic> baton;
         kv->asyncMultiPut(0, partId, std::move(data), [&](kvstore::ResultCode code) {
             EXPECT_EQ(code, kvstore::ResultCode::SUCCEEDED);
+            baton.post();
         });
+        baton.wait();
     }
-    sleep(1);
 }
 
 
diff --git a/src/storage/test/QueryVertexPropsTest.cpp b/src/storage/test/QueryVertexPropsTest.cpp
index 87c15ca4412..1bf585d9cd9 100644
--- a/src/storage/test/QueryVertexPropsTest.cpp
+++ b/src/storage/test/QueryVertexPropsTest.cpp
@@ -41,16 +41,17 @@ TEST(QueryVertexPropsTest, SimpleTest) {
                 data.emplace_back(std::move(key), std::move(val));
             }
         }
+        folly::Baton<true, std::atomic> baton;
         kv->asyncMultiPut(
             0,
             partId,
             std::move(data),
             [&](kvstore::ResultCode code) {
                 EXPECT_EQ(code, kvstore::ResultCode::SUCCEEDED);
+                baton.post();
             });
+        baton.wait();
     }
-
-    sleep(1);
     LOG(INFO) << "Build VertexPropsRequest...";
     cpp2::VertexPropRequest req;
     req.set_space_id(0);
diff --git a/src/storage/test/TestUtils.h b/src/storage/test/TestUtils.h
index e79b3031c85..7e2c699b8d9 100644
--- a/src/storage/test/TestUtils.h
+++ b/src/storage/test/TestUtils.h
@@ -17,6 +17,7 @@
 #include "dataman/ResultSchemaProvider.h"
 #include "storage/StorageServiceHandler.h"
 #include <thrift/lib/cpp2/server/ThriftServer.h>
+#include <folly/synchronization/Baton.h>
 #include "meta/SchemaManager.h"
 #include <folly/executors/ThreadPoolExecutor.h>
 #include <folly/executors/CPUThreadPoolExecutor.h>