intermediate changes

ref:cde9a383711a11544ce7e107a78147fb96cc4029

1 files changed, 104 insertions, 0 deletions

diff --git a/library/cpp/testing/unittest/tests_data.cpp b/library/cpp/testing/unittest/tests_data.cpp
new file mode 100644
index 0000000000..b51cbc4b87
--- /dev/null
+++ b/library/cpp/testing/unittest/tests_data.cpp
@@ -0,0 +1,104 @@
+#include "tests_data.h"
+#include "registar.h"
+
+#include <library/cpp/testing/common/network.h>
+
+#include <util/system/env.h>
+#include <util/system/mutex.h>
+
+class TPortManager::TPortManagerImpl {
+public:
+    TPortManagerImpl(bool reservePortsForCurrentTest)
+        : EnableReservePortsForCurrentTest(reservePortsForCurrentTest)
+        , DisableRandomPorts(!GetEnv("NO_RANDOM_PORTS").empty())
+    {
+    }
+
+    ui16 GetPort(ui16 port) {
+        if (port && DisableRandomPorts) {
+            return port;
+        }
+
+        TAtomicSharedPtr<NTesting::IPort> holder(NTesting::GetFreePort().Release());
+        ReservePortForCurrentTest(holder);
+
+        TGuard<TMutex> g(Lock);
+        ReservedPorts.push_back(holder);
+        return holder->Get();
+    }
+
+    ui16 GetUdpPort(ui16 port) {
+        return GetPort(port);
+    }
+
+    ui16 GetTcpPort(ui16 port) {
+        return GetPort(port);
+    }
+
+    ui16 GetTcpAndUdpPort(ui16 port) {
+        return GetPort(port);
+    }
+
+    ui16 GetPortsRange(const ui16 startPort, const ui16 range) {
+        Y_UNUSED(startPort);
+        auto ports = NTesting::NLegacy::GetFreePortsRange(range);
+        ui16 first = ports[0];
+        TGuard<TMutex> g(Lock);
+        for (auto& port : ports) {
+            ReservedPorts.emplace_back(port.Release());
+            ReservePortForCurrentTest(ReservedPorts.back());
+        }
+        return first;
+    }
+
+private:
+    void ReservePortForCurrentTest(const TAtomicSharedPtr<NTesting::IPort>& portGuard) {
+        if (EnableReservePortsForCurrentTest) {
+            TTestBase* currentTest = NUnitTest::NPrivate::GetCurrentTest();
+            if (currentTest != nullptr) {
+                currentTest->RunAfterTest([guard = portGuard]() mutable {
+                    guard = nullptr; // remove reference for allocated port
+                });
+            }
+        }
+    }
+
+private:
+    TMutex Lock;
+    TVector<TAtomicSharedPtr<NTesting::IPort>> ReservedPorts;
+    const bool EnableReservePortsForCurrentTest;
+    const bool DisableRandomPorts;
+};
+
+TPortManager::TPortManager(bool reservePortsForCurrentTest)
+    : Impl_(new TPortManagerImpl(reservePortsForCurrentTest))
+{
+}
+
+TPortManager::~TPortManager() {
+}
+
+ui16 TPortManager::GetPort(ui16 port) {
+    return Impl_->GetTcpPort(port);
+}
+
+ui16 TPortManager::GetTcpPort(ui16 port) {
+    return Impl_->GetTcpPort(port);
+}
+
+ui16 TPortManager::GetUdpPort(ui16 port) {
+    return Impl_->GetUdpPort(port);
+}
+
+ui16 TPortManager::GetTcpAndUdpPort(ui16 port) {
+    return Impl_->GetTcpAndUdpPort(port);
+}
+
+ui16 TPortManager::GetPortsRange(const ui16 startPort, const ui16 range) {
+    return Impl_->GetPortsRange(startPort, range);
+}
+
+ui16 GetRandomPort() {
+    TPortManager* pm = Singleton<TPortManager>(false);
+    return pm->GetPort();
+}