1# -*- coding: utf-8 -*-
2
3import os
4import unittest
5from StringIO import StringIO
6import textwrap
7
8import stringtemplate3
9
10from antlr3.dottreegen import toDOT
11from antlr3.treewizard import TreeWizard
12from antlr3.tree import CommonTreeAdaptor
13
14
15class TestToDOT(unittest.TestCase):
16    """Test case for the toDOT function."""
17
18    def setUp(self):
19        self.adaptor = CommonTreeAdaptor()
20        self.tokens = [
21            "", "", "", "", "", "A", "B", "C", "D", "E", "ID", "VAR"
22            ]
23        self.wiz = TreeWizard(self.adaptor, self.tokens)
24
25
26    def testNone(self):
27        """toDOT()"""
28
29        treeST = stringtemplate3.StringTemplate(
30            template=(
31            "digraph {\n" +
32            "  $nodes$\n" +
33            "  $edges$\n" +
34            "}\n")
35            )
36
37        edgeST = stringtemplate3.StringTemplate(
38            template="$parent$ -> $child$\n"
39            )
40
41        tree = self.wiz.create("(A B (B C C) (B (C D D)))")
42        st = toDOT(tree, self.adaptor, treeST, edgeST)
43
44        result = st.toString()
45        expected = textwrap.dedent(
46            '''\
47            digraph {
48              n0 [label="A"];
49              n1 [label="B"];
50              n2 [label="B"];
51              n3 [label="C"];
52              n4 [label="C"];
53              n5 [label="B"];
54              n6 [label="C"];
55              n7 [label="D"];
56              n8 [label="D"];
57
58              n0 -> n1
59              n0 -> n2
60              n2 -> n3
61              n2 -> n4
62              n0 -> n5
63              n5 -> n6
64              n6 -> n7
65              n6 -> n8
66
67            }
68            '''
69            )
70        self.assertEqual(result, expected)
71
72
73if __name__ == "__main__":
74    unittest.main(testRunner=unittest.TextTestRunner(verbosity=2))
75