xref: /aosp_15_r20/external/pytorch/test/cpp/c10d/StoreTestCommon.hpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
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 
set(Store & store,const std::string & key,const std::string & value)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 
compareSet(Store & store,const std::string & key,const std::string & expectedValue,const std::string & desiredValue)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 
check(Store & store,const std::string & key,const std::string & expected)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 
deleteKey(Store & store,const std::string & key)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