Home | History | Annotate | Download | only in altgraph_tests 
      1 import unittest
      2 import os
      3 
      4 from altgraph import Dot
      5 from altgraph import Graph
      6 from altgraph import GraphError
      7 
      8 
      9 class 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 
    369 if __name__ == "__main__": # pragma: no cover
    370     unittest.main()
    371