Graph generators for testing

infra/graph_generators.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Sun Mar 17 13:00:33 2019
+
+@author: Zargham
+"""
+import networkx as nx
+
+def lineGraphGen(N, bidir=False,nodeTypeName='vanilla',edgeTypeName='vanilla'):
+
+    line = nx.path_graph(N, create_using=nx.MultiDiGraph)
+    if not(bidir):
+        G = line
+    else:  
+        edges = line.edges
+        G = nx.MultiDiGraph()
+        for e in edges:
+            G.add_edge(e[0],e[1])
+            G.add_edge(e[1],e[0])
+
+    nx.set_node_attributes(G,nodeTypeName, 'type')
+    nx.set_edge_attributes(G,edgeTypeName, 'type')
+
+    return G
+
+def starGraphGen(N, kind='sink',nodeTypeName='vanilla',edgeTypeName='vanilla'):
+
+    star = nx.star_graph(N)
+    G = nx.MultiDiGraph()
+
+    for e in star.edges:
+        if (kind == 'source') or (kind == 'bidir'):
+            G.add_edge(e[0],e[1])
+        if (kind == 'sink') or (kind == 'bidir'):
+            G.add_edge(e[1],e[0])
+
+    nx.set_node_attributes(G,nodeTypeName, 'type')
+    nx.set_edge_attributes(G,edgeTypeName, 'type')
+
+    return G
+
+def circleGraphGen(N, bidir=False,nodeTypeName='vanilla',edgeTypeName='vanilla' ):
+
+    circle = nx.cycle_graph(N, create_using=nx.MultiDiGraph)
+    if not(bidir):
+        G = circle
+    else:  
+        edges = circle.edges
+        G = nx.MultiDiGraph()
+        for e in edges:
+            G.add_edge(e[0],e[1])
+            G.add_edge(e[1],e[0])
+
+    nx.set_node_attributes(G,nodeTypeName, 'type')
+    nx.set_edge_attributes(G,edgeTypeName, 'type')
+
+    return G
+
+def treeGraphGen(r,h, kind='sink',nodeTypeName='vanilla',edgeTypeName='vanilla'):
+
+    tree = nx.balanced_tree(r,h, create_using=nx.MultiDiGraph)
+
+    if kind=='source':
+        G = tree
+    elif kind =='sink':
+        G = nx.MultiDiGraph()
+        for e in tree.edges:
+            G.add_edge(e[1],e[0])
+    elif kind == 'bidir':
+        G = nx.MultiDiGraph()
+        for e in tree.edges:
+            G.add_edge(e[1],e[0])
+            G.add_edge(e[0],e[1])
+
+    nx.set_node_attributes(G,nodeTypeName, 'type')
+    nx.set_edge_attributes(G,edgeTypeName, 'type')
+
+    return G    

infra/graph_generators_test.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Mar 20 17:17:54 2019
+
+@author: Zargham
+"""
+import networkx as nx
+import unittest
+from graph_generators import lineGraphGen, starGraphGen, treeGraphGen
+
+class GraphGenerators_LineGraph(unittest.TestCase):
+    def test_lineGraphGen_sanity(self):
+        g = lineGraphGen(4)
+        self.assertEqual(list(g.edges()), [(0, 1), (1, 2), (2, 3)])
+
+    def test_lineGraphGen_bidir(self):
+        g = lineGraphGen(4, bidir=True)
+        # (node, neighbor)
+        self.assertEqual(list(g.edges()), [(0, 1), (1, 0), (1, 2), (2, 1), (2, 3), (3, 2)])
+
+class GraphGenerators_StarGraph(unittest.TestCase):
+
+# how to validate that 'kind' is either 'source', 'sink', or 'bidir' ?? 
+
+    def test_starGraph_sink(self):
+        g = starGraphGen(4)
+        self.assertEqual(list(g.edges()), [(1, 0), (2, 0), (3, 0), (4, 0)])
+
+    def test_starGraph_source(self):
+        g = starGraphGen(4, kind='source')
+        self.assertEqual(list(g.edges()), [(0, 1), (0, 2), (0, 3), (0, 4)])
+
+    def test_starGraph_bidir(self):
+        g = starGraphGen(4, kind='bidir')
+        self.assertEqual(list(g.edges()), [(0, 1), (0, 2), (0, 3), (0, 4), (1, 0), (2, 0), (3, 0), (4, 0)])
+
+class GraphGenerators_TreeGraph(unittest.TestCase):
+    #note: our default is diff from networkx's default in this case (sink vs source)
+    def test_treeGraph_source_sink(self):
+        g = treeGraphGen(2,2)
+        self.assertEqual(list(g.edges()), [(1, 0), (2, 0), (3, 1), (4, 1), (5, 2), (6, 2)])
+
+    def test_treeGraph_source(self):
+        g = treeGraphGen(2,2, kind='source')
+        self.assertEqual(list(g.edges()), [(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (2, 6)])
+
+    def test_treeGraph_bidir(self):
+        g = treeGraphGen(2,2, kind='bidir')
+        self.assertEqual(list(g.edges()), [
+            (1, 0), (1, 3), (1, 4), (0, 1), (0, 2), (2, 0),
+            (2, 5), (2, 6), (3, 1), (4, 1), (5, 2), (6, 2)
+        ])
+
+# this is the way I might test that the types assigned in Z's functions are correct 
+# I think it would be a good idea to abstract out that logic into a helper function 
+# so that the code isn't duplicated in every generator 
+
+class GraphGenerators_SetType(unittest.TestCase):
+    def test_lineGraphGen_setType(self):
+        g = lineGraphGen(4, nodeTypeName='foo',edgeTypeName='bar')
+
+        # each node and edge should have the correct 'type' value
+        for k, v in nx.get_node_attributes(g, 'type').items():
+            self.assertEqual(v, 'foo')
+
+        for k, v in nx.get_edge_attributes(g, 'type').items():
+            self.assertEqual(v, 'bar')