1import unittest
2import os
3
4from altgraph import Dot
5from altgraph import Graph
6from altgraph import GraphError
7
8
9class TestDot (unittest.TestCase):
10
11    def test_constructor(self):
12        g = Graph.Graph([
13                (1,2),
14                (1,3),
15                (1,4),
16                (2,4),
17                (2,6),
18                (2,7),
19                (7,4),
20                (6,1),
21            ]
22        )
23
24        dot = Dot.Dot(g)
25
26        self.assertEqual(dot.name, 'G')
27        self.assertEqual(dot.attr, {})
28        self.assertEqual(dot.temp_dot, 'tmp_dot.dot')
29        self.assertEqual(dot.temp_neo, 'tmp_neo.dot')
30        self.assertEqual(dot.dot, 'dot')
31        self.assertEqual(dot.dotty, 'dotty')
32        self.assertEqual(dot.neato, 'neato')
33        self.assertEqual(dot.type, 'digraph')
34
35        self.assertEqual(dot.nodes, dict([(x, {}) for x in g]))
36
37        edges = {}
38        for head in g:
39            edges[head] = {}
40            for tail in g.out_nbrs(head):
41                edges[head][tail] = {}
42
43        self.assertEqual(dot.edges[1], edges[1])
44        self.assertEqual(dot.edges, edges)
45
46
47        dot = Dot.Dot(g, nodes=[1,2],
48                edgefn=lambda node: list(sorted(g.out_nbrs(node)))[:-1],
49                nodevisitor=lambda node: {'label': node},
50                edgevisitor=lambda head, tail: {'label': (head, tail) },
51                name="testgraph",
52                dot='/usr/local/bin/dot',
53                dotty='/usr/local/bin/dotty',
54                neato='/usr/local/bin/neato',
55                graphtype="graph")
56
57        self.assertEqual(dot.name, 'testgraph')
58        self.assertEqual(dot.attr, {})
59        self.assertEqual(dot.temp_dot, 'tmp_dot.dot')
60        self.assertEqual(dot.temp_neo, 'tmp_neo.dot')
61        self.assertEqual(dot.dot, '/usr/local/bin/dot')
62        self.assertEqual(dot.dotty, '/usr/local/bin/dotty')
63        self.assertEqual(dot.neato, '/usr/local/bin/neato')
64        self.assertEqual(dot.type, 'graph')
65
66        self.assertEqual(dot.nodes, dict([(x, {'label': x}) for x in [1,2]]))
67
68        edges = {}
69        for head in [1,2]:
70            edges[head] = {}
71            for tail in list(sorted(g.out_nbrs(head)))[:-1]:
72                if tail not in [1,2]: continue
73                edges[head][tail] = {'label': (head, tail) }
74
75        self.assertEqual(dot.edges[1], edges[1])
76        self.assertEqual(dot.edges, edges)
77
78        self.assertRaises(GraphError, Dot.Dot, g, nodes=[1,2, 9])
79
80    def test_style(self):
81        g = Graph.Graph([])
82
83        dot = Dot.Dot(g)
84
85        self.assertEqual(dot.attr, {})
86
87        dot.style(key='value')
88        self.assertEqual(dot.attr, {'key': 'value'})
89
90        dot.style(key2='value2')
91        self.assertEqual(dot.attr, {'key2': 'value2'})
92
93    def test_node_style(self):
94        g = Graph.Graph([
95                (1,2),
96                (1,3),
97                (1,4),
98                (2,4),
99                (2,6),
100                (2,7),
101                (7,4),
102                (6,1),
103            ]
104        )
105
106        dot = Dot.Dot(g)
107
108        self.assertEqual(dot.nodes[1], {})
109
110        dot.node_style(1, key='value')
111        self.assertEqual(dot.nodes[1], {'key': 'value'})
112
113        dot.node_style(1, key2='value2')
114        self.assertEqual(dot.nodes[1], {'key2': 'value2'})
115        self.assertEqual(dot.nodes[2], {})
116
117        dot.all_node_style(key3='value3')
118        for n in g:
119            self.assertEqual(dot.nodes[n], {'key3': 'value3'})
120
121        self.assertTrue(9 not in dot.nodes)
122        dot.node_style(9, key='value')
123        self.assertEqual(dot.nodes[9], {'key': 'value'})
124
125    def test_edge_style(self):
126        g = Graph.Graph([
127                (1,2),
128                (1,3),
129                (1,4),
130                (2,4),
131                (2,6),
132                (2,7),
133                (7,4),
134                (6,1),
135            ]
136        )
137
138        dot = Dot.Dot(g)
139
140        self.assertEqual(dot.edges[1][2], {})
141        dot.edge_style(1,2, foo='bar')
142        self.assertEqual(dot.edges[1][2], {'foo': 'bar'})
143
144        dot.edge_style(1,2, foo2='2bar')
145        self.assertEqual(dot.edges[1][2], {'foo2': '2bar'})
146
147        self.assertEqual(dot.edges[1][3], {})
148
149        self.assertFalse(6 in dot.edges[1])
150        dot.edge_style(1,6, foo2='2bar')
151        self.assertEqual(dot.edges[1][6], {'foo2': '2bar'})
152
153        self.assertRaises(GraphError, dot.edge_style, 1, 9, a=1)
154        self.assertRaises(GraphError, dot.edge_style, 9, 1, a=1)
155
156
157    def test_iter(self):
158        g = Graph.Graph([
159                (1,2),
160                (1,3),
161                (1,4),
162                (2,4),
163                (2,6),
164                (2,7),
165                (7,4),
166                (6,1),
167            ]
168        )
169
170        dot = Dot.Dot(g)
171        dot.style(graph="foobar")
172        dot.node_style(1, key='value')
173        dot.node_style(2, key='another', key2='world')
174        dot.edge_style(1,4, key1='value1', key2='value2')
175        dot.edge_style(2,4, key1='valueA')
176
177        self.assertEqual(list(iter(dot)), list(dot.iterdot()))
178
179        for item in dot.iterdot():
180            self.assertTrue(isinstance(item, str))
181
182        first = list(dot.iterdot())[0]
183        self.assertEqual(first, "digraph %s {\n"%(dot.name,))
184
185        dot.type = 'graph'
186        first = list(dot.iterdot())[0]
187        self.assertEqual(first, "graph %s {\n"%(dot.name,))
188
189        dot.type = 'foo'
190        self.assertRaises(GraphError, list, dot.iterdot())
191        dot.type = 'digraph'
192
193        self.assertEqual(list(dot), [
194            'digraph G {\n',
195              'graph="foobar";',
196              '\n',
197
198            '\t"1" [',
199              'key="value",',
200            '];\n',
201
202            '\t"2" [',
203              'key="another",',
204              'key2="world",',
205            '];\n',
206
207            '\t"3" [',
208            '];\n',
209
210            '\t"4" [',
211            '];\n',
212
213            '\t"6" [',
214            '];\n',
215
216            '\t"7" [',
217            '];\n',
218
219            '\t"1" -> "2" [',
220            '];\n',
221
222            '\t"1" -> "3" [',
223            '];\n',
224
225            '\t"1" -> "4" [',
226              'key1="value1",',
227              'key2="value2",',
228            '];\n',
229
230             '\t"2" -> "4" [',
231               'key1="valueA",',
232             '];\n',
233
234             '\t"2" -> "6" [',
235             '];\n',
236
237             '\t"2" -> "7" [',
238             '];\n',
239
240             '\t"6" -> "1" [',
241             '];\n',
242
243             '\t"7" -> "4" [',
244             '];\n',
245           '}\n'])
246
247
248    def test_save(self):
249        g = Graph.Graph([
250                (1,2),
251                (1,3),
252                (1,4),
253                (2,4),
254                (2,6),
255                (2,7),
256                (7,4),
257                (6,1),
258            ]
259        )
260
261        dot = Dot.Dot(g)
262        dot.style(graph="foobar")
263        dot.node_style(1, key='value')
264        dot.node_style(2, key='another', key2='world')
265        dot.edge_style(1,4, key1='value1', key2='value2')
266        dot.edge_style(2,4, key1='valueA')
267
268        fn = 'test_dot.dot'
269        self.assertTrue(not os.path.exists(fn))
270
271        try:
272            dot.save_dot(fn)
273
274            fp = open(fn, 'r')
275            data = fp.read()
276            fp.close()
277            self.assertEqual(data, ''.join(dot))
278
279        finally:
280            if os.path.exists(fn):
281                os.unlink(fn)
282
283
284    def test_img(self):
285        g = Graph.Graph([
286                (1,2),
287                (1,3),
288                (1,4),
289                (2,4),
290                (2,6),
291                (2,7),
292                (7,4),
293                (6,1),
294            ]
295        )
296
297        dot = Dot.Dot(g, dot='/usr/local/bin/!!dot', dotty='/usr/local/bin/!!dotty', neato='/usr/local/bin/!!neato')
298        dot.style(size='10,10', rankdir='RL', page='5, 5' , ranksep=0.75)
299        dot.node_style(1, label='BASE_NODE',shape='box', color='blue')
300        dot.node_style(2, style='filled', fillcolor='red')
301        dot.edge_style(1,4, style='dotted')
302        dot.edge_style(2,4, arrowhead='dot', label='binds', labelangle='90')
303
304        system_cmds = []
305        def fake_system(cmd):
306            system_cmds.append(cmd)
307            return None
308
309        try:
310            real_system = os.system
311            os.system = fake_system
312
313            system_cmds = []
314            dot.save_img('foo')
315            self.assertEqual(system_cmds, ['/usr/local/bin/!!dot -Tgif tmp_dot.dot -o foo.gif'])
316
317            system_cmds = []
318            dot.save_img('foo', file_type='jpg')
319            self.assertEqual(system_cmds, ['/usr/local/bin/!!dot -Tjpg tmp_dot.dot -o foo.jpg'])
320
321            system_cmds = []
322            dot.save_img('bar', file_type='jpg', mode='neato')
323            self.assertEqual(system_cmds, [
324                '/usr/local/bin/!!neato -o tmp_dot.dot tmp_neo.dot',
325                '/usr/local/bin/!!dot -Tjpg tmp_dot.dot -o bar.jpg',
326            ])
327
328            system_cmds = []
329            dot.display()
330            self.assertEqual(system_cmds, [
331                '/usr/local/bin/!!dotty tmp_dot.dot'
332            ])
333
334            system_cmds = []
335            dot.display(mode='neato')
336            self.assertEqual(system_cmds, [
337                '/usr/local/bin/!!neato -o tmp_dot.dot tmp_neo.dot',
338                '/usr/local/bin/!!dotty tmp_dot.dot'
339            ])
340
341        finally:
342            if os.path.exists(dot.temp_dot):
343                os.unlink(dot.temp_dot)
344            if os.path.exists(dot.temp_neo):
345                os.unlink(dot.temp_neo)
346            os.system = real_system
347
348        if os.path.exists('/usr/local/bin/dot') and os.path.exists('/usr/local/bin/neato'):
349            try:
350                dot.dot='/usr/local/bin/dot'
351                dot.neato='/usr/local/bin/neato'
352                self.assertFalse(os.path.exists('foo.gif'))
353                dot.save_img('foo')
354                self.assertTrue(os.path.exists('foo.gif'))
355                os.unlink('foo.gif')
356
357                self.assertFalse(os.path.exists('foo.gif'))
358                dot.save_img('foo', mode='neato')
359                self.assertTrue(os.path.exists('foo.gif'))
360                os.unlink('foo.gif')
361
362            finally:
363                if os.path.exists(dot.temp_dot):
364                    os.unlink(dot.temp_dot)
365                if os.path.exists(dot.temp_neo):
366                    os.unlink(dot.temp_neo)
367
368
369if __name__ == "__main__": # pragma: no cover
370    unittest.main()
371