Add gtests for MPPTunnel (#4553)

close #4481
pingcap · Apr 18, 2022 · 605ddc0 · 605ddc0
1 parent 66f45c7
commit 605ddc0
Show file tree

Hide file tree

Showing 3 changed files with 669 additions and 0 deletions.
diff --git a/dbms/src/Flash/Mpp/MPPTunnel.cpp b/dbms/src/Flash/Mpp/MPPTunnel.cpp
@@ -51,6 +51,28 @@ MPPTunnelBase<Writer>::MPPTunnelBase(
     GET_METRIC(tiflash_object_count, type_count_of_mpptunnel).Increment();
 }
 
+template <typename Writer>
+MPPTunnelBase<Writer>::MPPTunnelBase(
+    const String & tunnel_id_,
+    const std::chrono::seconds timeout_,
+    int input_steams_num_,
+    bool is_local_,
+    bool is_async_,
+    const String & req_id)
+    : connected(false)
+    , finished(false)
+    , is_local(is_local_)
+    , is_async(is_async_)
+    , timeout(timeout_)
+    , tunnel_id(tunnel_id_)
+    , input_streams_num(input_steams_num_)
+    , send_queue(std::max(5, input_steams_num_ * 5)) // MPMCQueue can benefit from a slightly larger queue size
+    , thread_manager(newThreadManager())
+    , log(Logger::get("MPPTunnel", req_id, tunnel_id))
+{
+    RUNTIME_ASSERT(!(is_local && is_async), log, "is_local: {}, is_async: {}.", is_local, is_async);
+}
+
 template <typename Writer>
 MPPTunnelBase<Writer>::~MPPTunnelBase()
 {

diff --git a/dbms/src/Flash/Mpp/MPPTunnel.h b/dbms/src/Flash/Mpp/MPPTunnel.h
@@ -42,6 +42,12 @@
 
 namespace DB
 {
+namespace tests
+{
+class MPPTunnelTest;
+class TestMPPTunnelBase;
+} // namespace tests
+
 class EstablishCallData;
 
 /**
@@ -123,6 +129,17 @@ class MPPTunnelBase : private boost::noncopyable
     void sendJob(bool need_lock = true);
 
 private:
+    friend class tests::MPPTunnelTest;
+    friend class tests::TestMPPTunnelBase;
+    // For gtest usage
+    MPPTunnelBase(
+        const String & tunnel_id_,
+        std::chrono::seconds timeout_,
+        int input_steams_num_,
+        bool is_local_,
+        bool is_async_,
+        const String & req_id);
+
     void finishSendQueue();
 
     void waitUntilConnectedOrFinished(std::unique_lock<std::mutex> & lk);