1 | #pragma once |
2 | |
3 | #include <torch/csrc/distributed/c10d/Store.hpp> |
4 | #include "TestUtils.hpp" |
5 | |
6 | #include <gtest/gtest.h> |
7 | |
8 | namespace c10d { |
9 | namespace test { |
10 | |
11 | inline void set( |
12 | Store& store, |
13 | const std::string& key, |
14 | const std::string& value) { |
15 | std::vector<uint8_t> data(value.begin(), value.end()); |
16 | store.set(key, data); |
17 | } |
18 | |
19 | inline std::vector<uint8_t> compareSet( |
20 | Store& store, |
21 | const std::string& key, |
22 | const std::string& expectedValue, |
23 | const std::string& desiredValue) { |
24 | std::vector<uint8_t> expectedData(expectedValue.begin(), expectedValue.end()); |
25 | std::vector<uint8_t> desiredData(desiredValue.begin(), desiredValue.end()); |
26 | return store.compareSet(key, expectedData, desiredData); |
27 | } |
28 | |
29 | inline void check( |
30 | Store& store, |
31 | const std::string& key, |
32 | const std::string& expected) { |
33 | auto tmp = store.get(key); |
34 | auto actual = std::string((const char*)tmp.data(), tmp.size()); |
35 | EXPECT_EQ(actual, expected); |
36 | } |
37 | |
38 | inline void deleteKey( |
39 | Store& store, |
40 | const std::string& key) { |
41 | store.deleteKey(key); |
42 | } |
43 | |
44 | } // namespace test |
45 | } // namespace c10d |
46 | |