def test_graph_graphviz_source_pstate(self): """ test the output of graphviz source, with the `pstate_node_styles` function """ nodes = self.create_provenance() graph = graph_mod.Graph(node_style_fn=graph_mod.pstate_node_styles) graph.recurse_descendants(nodes.pd0) # print() # print(graph.graphviz.source) # graph.graphviz.render("test_graphviz_pstate", cleanup=True) expected = """\ digraph {{ N{pd0} [label="Dict ({pd0})" pencolor=black shape=rectangle] N{calc1} [label="CalcJobNode ({calc1}) State: finished Exit Code: 0" fillcolor="#8cd499ff" penwidth=0 shape=ellipse style=filled] N{pd0} -> N{calc1} [color="#000000" style=solid] N{wc1} [label="WorkChainNode ({wc1}) State: running" fillcolor="#e38851ff" penwidth=0 shape=polygon sides=6 style=filled] N{pd0} -> N{wc1} [color="#000000" style=dashed] N{rd1} [label="RemoteData ({rd1}) @localhost" pencolor=black shape=rectangle] N{calc1} -> N{rd1} [color="#000000" style=solid] N{fd1} [label="FolderData ({fd1})" pencolor=black shape=rectangle] N{wc1} -> N{fd1} [color="#000000" style=dashed] N{pd3} [label="Dict ({pd3})" pencolor=black shape=rectangle] N{wc1} -> N{pd3} [color="#000000" style=dashed] N{calcf1} [label="CalcFunctionNode ({calcf1}) State: finished Exit Code: 200" fillcolor="#de707fff" penwidth=0 shape=ellipse style=filled] N{wc1} -> N{calcf1} [color="#000000" style=dotted] N{wc1} -> N{calc1} [color="#000000" style=dotted] N{rd1} -> N{calcf1} [color="#000000" style=solid] N{calcf1} -> N{fd1} [color="#000000" style=solid] N{calcf1} -> N{pd3} [color="#000000" style=solid] }}""".format(**{k: v.pk for k, v in nodes.items()}) # dedent before comparison self.assertEqual( sorted([l.strip() for l in graph.graphviz.source.splitlines()]), sorted([l.strip() for l in expected.splitlines()]))
def test_graph_recurse_descendants(self): """ test adding nodes and all its (recursed) incoming nodes to a graph""" nodes = self.create_provenance() graph = graph_mod.Graph() graph.recurse_descendants(nodes.rd1) self.assertEqual( graph.nodes, set([nodes.rd1.pk, nodes.calcf1.pk, nodes.pd3.pk, nodes.fd1.pk])) self.assertEqual( graph.edges, set([ (nodes.rd1.pk, nodes.calcf1.pk, LinkPair(LinkType.INPUT_CALC, 'input1')), (nodes.calcf1.pk, nodes.pd3.pk, LinkPair(LinkType.CREATE, 'output1')), (nodes.calcf1.pk, nodes.fd1.pk, LinkPair(LinkType.CREATE, 'output2')), ]))
def test_graph_init(self): """ test initialisation of Graph""" # pylint: disable=no-self-use graph_mod.Graph()