xref: /aosp_15_r20/external/pytorch/test/cpp/lazy/test_util.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #include <gtest/gtest.h>
2 
3 #include <exception>
4 
5 #include <torch/csrc/lazy/core/util.h>
6 
7 namespace torch {
8 namespace lazy {
9 
TEST(UtilTest,ExceptionCleanup)10 TEST(UtilTest, ExceptionCleanup) {
11   std::exception_ptr exception;
12   EXPECT_EQ(exception, nullptr);
13 
14   {
15     ExceptionCleanup cleanup(
16         [&](std::exception_ptr&& e) { exception = std::move(e); });
17 
18     cleanup.SetStatus(std::make_exception_ptr(std::runtime_error("Oops!")));
19   }
20   EXPECT_NE(exception, nullptr);
21 
22   try {
23     std::rethrow_exception(exception);
24   } catch (const std::exception& e) {
25     EXPECT_STREQ(e.what(), "Oops!");
26   }
27 
28   exception = nullptr;
29   {
30     ExceptionCleanup cleanup(
31         [&](std::exception_ptr&& e) { exception = std::move(e); });
32 
33     cleanup.SetStatus(std::make_exception_ptr(std::runtime_error("")));
34     cleanup.Release();
35   }
36   EXPECT_EQ(exception, nullptr);
37 }
38 
TEST(UtilTest,MaybeRef)39 TEST(UtilTest, MaybeRef) {
40   std::string storage("String storage");
41   MaybeRef<std::string> refStorage(storage);
42   EXPECT_FALSE(refStorage.IsStored());
43   EXPECT_EQ(*refStorage, storage);
44 
45   MaybeRef<std::string> effStorage(std::string("Vanishing"));
46   EXPECT_TRUE(effStorage.IsStored());
47   EXPECT_EQ(*effStorage, "Vanishing");
48 }
49 
TEST(UtilTest,Iota)50 TEST(UtilTest, Iota) {
51   auto result = Iota<int>(0);
52   EXPECT_TRUE(result.empty());
53 
54   result = Iota<int>(1);
55   EXPECT_EQ(result.size(), 1);
56   EXPECT_EQ(result[0], 0);
57 
58   result = Iota<int>(2);
59   EXPECT_EQ(result.size(), 2);
60   EXPECT_EQ(result[0], 0);
61   EXPECT_EQ(result[1], 1);
62 
63   result = Iota<int>(3, 1, 3);
64   EXPECT_EQ(result.size(), 3);
65   EXPECT_EQ(result[0], 1);
66   EXPECT_EQ(result[1], 4);
67   EXPECT_EQ(result[2], 7);
68 }
69 
70 } // namespace lazy
71 } // namespace torch
72