1 #include <gtest/gtest.h>
2
3 #include <ATen/code_template.h>
4 #include <test/cpp/jit/test_utils.h>
5
6 namespace torch {
7 namespace jit {
8
9 static const auto ct = at::jit::CodeTemplate(R"(
10 int foo($args) {
11
12 $bar
13 $bar
14 $a+$b
15 }
16 int commatest(int a${,stuff})
17 int notest(int a${,empty,})
18 )");
19 static const auto ct_expect = R"(
20 int foo(hi, 8) {
21
22 what
23 on many
24 lines...
25 7
26 what
27 on many
28 lines...
29 7
30 3+4
31 }
32 int commatest(int a, things..., others)
33 int notest(int a)
34 )";
35
TEST(TestCodeTemplate,Copying)36 TEST(TestCodeTemplate, Copying) {
37 at::jit::TemplateEnv e;
38 e.s("hi", "foo");
39 e.v("what", {"is", "this"});
40 at::jit::TemplateEnv c(e);
41 c.s("hi", "foo2");
42 ASSERT_EQ(e.s("hi"), "foo");
43 ASSERT_EQ(c.s("hi"), "foo2");
44 ASSERT_EQ(e.v("what")[0], "is");
45 }
46
TEST(TestCodeTemplate,Formatting)47 TEST(TestCodeTemplate, Formatting) {
48 at::jit::TemplateEnv e;
49 e.v("args", {"hi", "8"});
50 e.v("bar", {"what\non many\nlines...", "7"});
51 e.s("a", "3");
52 e.s("b", "4");
53 e.v("stuff", {"things...", "others"});
54 e.v("empty", {});
55 auto s = ct.format(e);
56 // std::cout << "'" << s << "'\n";
57 // std::cout << "'" << ct_expect << "'\n";
58 ASSERT_EQ(s, ct_expect);
59 }
60
61 } // namespace jit
62 } // namespace torch
63