Skip to content

Commit 75939d5

Browse files
committed
Add tests
1 parent 1286d20 commit 75939d5

File tree

1 file changed

+121
-13
lines changed

1 file changed

+121
-13
lines changed

tests/unit/client/local_dc_detector/local_dc_detector_ut.cpp

Lines changed: 121 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -42,14 +42,23 @@ class TMockedPinger {
4242

4343
TDuration operator()(const Ydb::Discovery::EndpointInfo& endpoint) const {
4444
auto it = EndpointByAdress_.find(endpoint.address());
45-
if (it == EndpointByAdress_.end()) {
45+
if (it == EndpointByAdress_.end() || Blacklist_.contains(endpoint.address())) {
4646
return TDuration::Max();
4747
}
4848
return it->second.Ping();
4949
}
5050

51+
void BanEndpoint(const std::string& adress) {
52+
Blacklist_.insert(adress);
53+
}
54+
55+
void UnbanEndpoint(const std::string& adress) {
56+
Blacklist_.erase(adress);
57+
}
58+
5159
private:
5260
mutable std::unordered_map<std::string, TMockedEndpoint> EndpointByAdress_;
61+
std::unordered_set<std::string> Blacklist_;
5362
};
5463

5564
std::vector<TDuration> GenerateMeasures(size_t count, int minMs, int maxMs, std::mt19937& gen) {
@@ -62,7 +71,6 @@ std::vector<TDuration> GenerateMeasures(size_t count, int minMs, int maxMs, std:
6271
return measures;
6372
}
6473

65-
6674
Y_UNIT_TEST_SUITE(LocalDCDetectionTest) {
6775
Y_UNIT_TEST(Basic) {
6876
Ydb::Discovery::ListEndpointsResult endpoints;
@@ -73,20 +81,23 @@ Y_UNIT_TEST_SUITE(LocalDCDetectionTest) {
7381
const std::vector<std::string> endpointsB = {"B1", "B2", "B3"};
7482
const std::vector<std::string> endpointsC = {"C1", "C2", "C3", "C4", "C5", "C6", "C7"};
7583

84+
const std::size_t epoches = 3;
85+
const std::size_t measuresAmount = 10 * epoches;
86+
7687
for (const auto& ep : endpointsA) {
77-
mockData[ep] = GenerateMeasures(10, 20, 30, gen);
88+
mockData[ep] = GenerateMeasures(measuresAmount, 20, 30, gen);
7889
auto& endpoint = *endpoints.add_endpoints();
7990
endpoint.set_location("A");
8091
endpoint.set_address(ep);
8192
}
8293
for (const auto& ep : endpointsB) {
83-
mockData[ep] = GenerateMeasures(10, 30, 45, gen);
94+
mockData[ep] = GenerateMeasures(measuresAmount, 30, 45, gen);
8495
auto& endpoint = *endpoints.add_endpoints();
8596
endpoint.set_location("B");
8697
endpoint.set_address(ep);
8798
}
8899
for (const auto& ep : endpointsC) {
89-
mockData[ep] = GenerateMeasures(8, 50, 70, gen);
100+
mockData[ep] = GenerateMeasures(measuresAmount, 50, 70, gen);
90101
auto& endpoint = *endpoints.add_endpoints();
91102
endpoint.set_location("C");
92103
endpoint.set_address(ep);
@@ -95,44 +106,141 @@ Y_UNIT_TEST_SUITE(LocalDCDetectionTest) {
95106
std::function<TDuration(const Ydb::Discovery::EndpointInfo& endpoint)> pinger = TMockedPinger(mockData);
96107
TLocalDCDetector detector(pinger);
97108

109+
for (std::size_t i = 0; i < epoches; ++i) {
110+
detector.DetectLocalDC(endpoints);
111+
UNIT_ASSERT_VALUES_EQUAL(true, detector.IsLocalDC("A"));
112+
UNIT_ASSERT_VALUES_EQUAL(false, detector.IsLocalDC("B"));
113+
UNIT_ASSERT_VALUES_EQUAL(false, detector.IsLocalDC("C"));
114+
}
115+
}
116+
117+
Y_UNIT_TEST(SingleLocation) {
118+
Ydb::Discovery::ListEndpointsResult endpoints;
119+
std::unordered_map<std::string, std::vector<TDuration>> mockData;
120+
std::mt19937 gen(std::random_device{}());
121+
122+
const std::vector<std::string> endpointsA = {"A1", "A2", "A3"};
123+
124+
for (const auto& ep : endpointsA) {
125+
mockData[ep] = GenerateMeasures(10, 20, 30, gen);
126+
auto& endpoint = *endpoints.add_endpoints();
127+
endpoint.set_location("A");
128+
endpoint.set_address(ep);
129+
}
130+
131+
std::function<TDuration(const Ydb::Discovery::EndpointInfo& endpoint)> pinger = TMockedPinger(mockData);
132+
TLocalDCDetector detector(pinger);
133+
98134
detector.DetectLocalDC(endpoints);
99135

100-
UNIT_ASSERT_VALUES_EQUAL(false, detector.IsLocalDC("C"));
101-
UNIT_ASSERT_VALUES_EQUAL(false, detector.IsLocalDC("B"));
102136
UNIT_ASSERT_VALUES_EQUAL(true, detector.IsLocalDC("A"));
103137
}
104138

105-
Y_UNIT_TEST(Fallback) {
139+
Y_UNIT_TEST(UnavailableLocalDC) {
106140
Ydb::Discovery::ListEndpointsResult endpoints;
107141
std::unordered_map<std::string, std::vector<TDuration>> mockData;
142+
std::mt19937 gen(std::random_device{}());
108143

109-
const std::vector<std::string> endpointsA = {"A1", "A2", "A3"};
110-
const std::vector<std::string> endpointsB = {"B1", "B2", "B3"};
111-
const std::vector<std::string> endpointsC = {"C1", "C2", "C3"};
144+
const std::vector<std::string> endpointsA = {"A1", "A2", "A3", "A4", "A5", "A6", "A7"};
145+
const std::vector<std::string> endpointsB = {"B1", "B2", "B3", "B4", "B5", "B6", "B7"};
146+
const std::vector<std::string> endpointsC = {"C1", "C2", "C3", "C4", "C5", "C6", "C7"};
147+
148+
const std::size_t epoches = 3;
149+
const std::size_t measuresAmount = 10 * epoches;
112150

113151
for (const auto& ep : endpointsA) {
152+
mockData[ep] = GenerateMeasures(measuresAmount, 20, 30, gen);
114153
auto& endpoint = *endpoints.add_endpoints();
115154
endpoint.set_location("A");
116155
endpoint.set_address(ep);
117156
}
118157
for (const auto& ep : endpointsB) {
158+
mockData[ep] = GenerateMeasures(measuresAmount, 30, 45, gen);
119159
auto& endpoint = *endpoints.add_endpoints();
120160
endpoint.set_location("B");
121161
endpoint.set_address(ep);
122162
}
123163
for (const auto& ep : endpointsC) {
164+
mockData[ep] = GenerateMeasures(measuresAmount, 50, 70, gen);
124165
auto& endpoint = *endpoints.add_endpoints();
125166
endpoint.set_location("C");
126167
endpoint.set_address(ep);
127168
}
128169

129-
std::function<TDuration(const Ydb::Discovery::EndpointInfo& endpoint)> pinger = TMockedPinger(mockData);
170+
TMockedPinger mockPinger(mockData);
171+
std::function<TDuration(const Ydb::Discovery::EndpointInfo& endpoint)> pinger =
172+
[&mockPinger](const Ydb::Discovery::EndpointInfo& endpoint) { return mockPinger(endpoint); };
173+
130174
TLocalDCDetector detector(pinger);
131175

132176
detector.DetectLocalDC(endpoints);
177+
UNIT_ASSERT_VALUES_EQUAL(true, detector.IsLocalDC("A"));
178+
UNIT_ASSERT_VALUES_EQUAL(false, detector.IsLocalDC("B"));
179+
UNIT_ASSERT_VALUES_EQUAL(false, detector.IsLocalDC("C"));
133180

134-
UNIT_ASSERT_VALUES_EQUAL(true, detector.IsLocalDC("C"));
181+
for (const auto& ep : endpointsA) {
182+
mockPinger.BanEndpoint(ep);
183+
}
184+
185+
detector.DetectLocalDC(endpoints);
186+
UNIT_ASSERT_VALUES_EQUAL(false, detector.IsLocalDC("A"));
135187
UNIT_ASSERT_VALUES_EQUAL(true, detector.IsLocalDC("B"));
188+
UNIT_ASSERT_VALUES_EQUAL(false, detector.IsLocalDC("C"));
189+
190+
for (const auto& ep : endpointsA) {
191+
mockPinger.UnbanEndpoint(ep);
192+
}
193+
194+
detector.DetectLocalDC(endpoints);
136195
UNIT_ASSERT_VALUES_EQUAL(true, detector.IsLocalDC("A"));
196+
UNIT_ASSERT_VALUES_EQUAL(false, detector.IsLocalDC("B"));
197+
UNIT_ASSERT_VALUES_EQUAL(false, detector.IsLocalDC("C"));
198+
}
199+
200+
Y_UNIT_TEST(OfflineDCs) {
201+
Ydb::Discovery::ListEndpointsResult endpoints;
202+
std::unordered_map<std::string, std::vector<TDuration>> mockData;
203+
std::mt19937 gen(std::random_device{}());
204+
205+
const std::vector<std::string> endpointsA = {"A1", "A2", "A3", "A4", "A5", "A6", "A7"};
206+
const std::vector<std::string> endpointsB = {"B1", "B2", "B3", "B4", "B5", "B6", "B7"};
207+
const std::vector<std::string> endpointsC = {"C1", "C2", "C3", "C4", "C5", "C6", "C7"};
208+
209+
for (const auto& ep : endpointsA) {
210+
auto& endpoint = *endpoints.add_endpoints();
211+
endpoint.set_location("A");
212+
endpoint.set_address(ep);
213+
}
214+
for (const auto& ep : endpointsB) {
215+
auto& endpoint = *endpoints.add_endpoints();
216+
endpoint.set_location("B");
217+
endpoint.set_address(ep);
218+
}
219+
for (const auto& ep : endpointsC) {
220+
auto& endpoint = *endpoints.add_endpoints();
221+
endpoint.set_location("C");
222+
endpoint.set_address(ep);
223+
}
224+
225+
TMockedPinger mockPinger(mockData);
226+
std::function<TDuration(const Ydb::Discovery::EndpointInfo& endpoint)> pinger =
227+
[&mockPinger](const Ydb::Discovery::EndpointInfo& endpoint) { return mockPinger(endpoint); };
228+
229+
TLocalDCDetector detector(pinger);
230+
231+
for (const auto& ep : endpointsA) {
232+
mockPinger.BanEndpoint(ep);
233+
}
234+
for (const auto& ep : endpointsB) {
235+
mockPinger.BanEndpoint(ep);
236+
}
237+
for (const auto& ep : endpointsC) {
238+
mockPinger.BanEndpoint(ep);
239+
}
240+
241+
detector.DetectLocalDC(endpoints);
242+
UNIT_ASSERT_VALUES_EQUAL(true, detector.IsLocalDC("A"));
243+
UNIT_ASSERT_VALUES_EQUAL(true, detector.IsLocalDC("B"));
244+
UNIT_ASSERT_VALUES_EQUAL(true, detector.IsLocalDC("C"));
137245
}
138246
}

0 commit comments

Comments
 (0)