1#include <gtest/gtest.h>
2
3#include <exception>
4
5#include <torch/csrc/lazy/core/util.h>
6
7namespace torch {
8namespace lazy {
9
10TEST(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
39TEST(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
50TEST(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