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()