diagrams/tests/test_diagram.py

import os
import shutil
import unittest

from diagrams import Cluster, Diagram, Node, Edge
from diagrams import getcluster, getdiagram, setcluster, setdiagram


class DiagramTest(unittest.TestCase):
    def setUp(self):
        self.name = "diagram_test"

    def tearDown(self):
        setdiagram(None)
        setcluster(None)
        # Only some tests generate the image file.
        try:
            shutil.rmtree(self.name)
        except OSError:
            # Consider it file
            try:
                os.remove(self.name + ".png")
            except FileNotFoundError:
                pass

    def test_validate_direction(self):
        # Normal directions.
        for dir in ("TB", "BT", "LR", "RL"):
            Diagram(direction=dir)

        # Invalid directions.
        for dir in ("BR", "TL", "Unknown"):
            with self.assertRaises(ValueError):
                Diagram(direction=dir)

    def test_validate_outformat(self):
        # Normal output formats.
        for fmt in ("png", "jpg", "svg", "pdf"):
            Diagram(outformat=fmt)

        # Invalid output formats.
        for fmt in ("pnp", "jpe", "unknown"):
            with self.assertRaises(ValueError):
                Diagram(outformat=fmt)

    def test_with_global_context(self):
        self.assertIsNone(getdiagram())
        with Diagram(name=os.path.join(self.name, 'with_global_context'), show=False):
            self.assertIsNotNone(getdiagram())
        self.assertIsNone(getdiagram())

    def test_node_not_in_diagram(self):
        # Node must be belong to a diagrams.
        with self.assertRaises(EnvironmentError):
            Node("node")

    def test_node_to_node(self):
        with Diagram(name=os.path.join(self.name, 'node_to_node'), show=False):
            node1 = Node("node1")
            node2 = Node("node2")
            self.assertEqual(node1 - node2, node2)
            self.assertEqual(node1 >> node2, node2)
            self.assertEqual(node1 << node2, node2)

    def test_node_to_nodes(self):
        with Diagram(name=os.path.join(self.name, 'node_to_nodes'), show=False):
            node1 = Node("node1")
            nodes = [Node("node2"), Node("node3")]
            self.assertEqual(node1 - nodes, nodes)
            self.assertEqual(node1 >> nodes, nodes)
            self.assertEqual(node1 << nodes, nodes)

    def test_nodes_to_node(self):
        with Diagram(name=os.path.join(self.name, 'nodes_to_node'), show=False):
            node1 = Node("node1")
            nodes = [Node("node2"), Node("node3")]
            self.assertEqual(nodes - node1, node1)
            self.assertEqual(nodes >> node1, node1)
            self.assertEqual(nodes << node1, node1)

    def test_default_filename(self):
        self.name = "example_1"
        with Diagram(name="Example 1", show=False):
            Node("node1")
        self.assertTrue(os.path.exists(f"{self.name}.png"))

    def test_custom_filename(self):
        self.name = "my_custom_name"
        with Diagram(name="Example 1", filename=self.name, show=False):
            Node("node1")
        self.assertTrue(os.path.exists(f"{self.name}.png"))


class ClusterTest(unittest.TestCase):
    def setUp(self):
        self.name = "cluster_test"

    def tearDown(self):
        setdiagram(None)
        setcluster(None)
        # Only some tests generate the image file.
        try:
            shutil.rmtree(self.name)
        except OSError:
            pass

    def test_validate_direction(self):
        # Normal directions.
        for dir in ("TB", "BT", "LR", "RL"):
            with Diagram(name=os.path.join(self.name, 'validate_direction'), show=False):
                Cluster(direction=dir)

        # Invalid directions.
        for dir in ("BR", "TL", "Unknown"):
            with self.assertRaises(ValueError):
                with Diagram(name=os.path.join(self.name, 'validate_direction'), show=False):
                    Cluster(direction=dir)

    def test_with_global_context(self):
        with Diagram(name=os.path.join(self.name, 'with_global_context'), show=False):
            self.assertIsNone(getcluster())
            with Cluster():
                self.assertIsNotNone(getcluster())
            self.assertIsNone(getcluster())

    def test_with_nested_cluster(self):
        with Diagram(name=os.path.join(self.name, 'with_nested_cluster'), show=False):
            self.assertIsNone(getcluster())
            with Cluster() as c1:
                self.assertEqual(c1, getcluster())
                with Cluster() as c2:
                    self.assertEqual(c2, getcluster())
                self.assertEqual(c1, getcluster())
            self.assertIsNone(getcluster())

    def test_node_not_in_diagram(self):
        # Node must be belong to a diagrams.
        with self.assertRaises(EnvironmentError):
            Node("node")

    def test_node_to_node(self):
        with Diagram(name=os.path.join(self.name, 'node_to_node'), show=False):
            with Cluster():
                node1 = Node("node1")
                node2 = Node("node2")
                self.assertEqual(node1 - node2, node2)
                self.assertEqual(node1 >> node2, node2)
                self.assertEqual(node1 << node2, node2)

    def test_node_to_nodes(self):
        with Diagram(name=os.path.join(self.name, 'node_to_nodes'), show=False):
            with Cluster():
                node1 = Node("node1")
                nodes = [Node("node2"), Node("node3")]
                self.assertEqual(node1 - nodes, nodes)
                self.assertEqual(node1 >> nodes, nodes)
                self.assertEqual(node1 << nodes, nodes)

    def test_nodes_to_node(self):
        with Diagram(name=os.path.join(self.name, 'nodes_to_node'), show=False):
            with Cluster():
                node1 = Node("node1")
                nodes = [Node("node2"), Node("node3")]
                self.assertEqual(nodes - node1, node1)
                self.assertEqual(nodes >> node1, node1)
                self.assertEqual(nodes << node1, node1)


class EdgeTest(unittest.TestCase):
    def setUp(self):
        self.name = "edge_test"

    def tearDown(self):
        setdiagram(None)
        setcluster(None)
        # Only some tests generate the image file.
        try:
            shutil.rmtree(self.name)
        except OSError:
            pass

    def test_node_to_node(self):
        with Diagram(name=os.path.join(self.name, 'node_to_node'), show=False):
            node1 = Node("node1")
            node2 = Node("node2")
            self.assertEqual(node1 - Edge(color='red') - node2, node2)

    def test_node_to_nodes(self):
        with Diagram(name=os.path.join(self.name, 'node_to_nodes'), show=False):
            with Cluster():
                node1 = Node("node1")
                nodes = [Node("node2"), Node("node3")]
                self.assertEqual(node1 - Edge(color='red') - nodes, nodes)

    def test_nodes_to_node(self):
        with Diagram(name=os.path.join(self.name, 'nodes_to_node'), show=False):
            with Cluster():
                node1 = Node("node1")
                nodes = [Node("node2"), Node("node3")]
                self.assertEqual(nodes - Edge(color='red') - node1, node1)

    def test_nodes_to_node_with_additional_attributes(self):
        with Diagram(name=os.path.join(self.name, 'nodes_to_node_with_additional_attributes'), show=False):
            with Cluster():
                node1 = Node("node1")
                nodes = [Node("node2"), Node("node3")]
                self.assertEqual(nodes - Edge(color='red') - Edge(color='green') - node1, node1)

    def test_node_to_node_with_attributes(self):
        with Diagram(name=os.path.join(self.name, 'node_to_node_with_attributes'), show=False):
            with Cluster():
                node1 = Node("node1")
                node2 = Node("node2")
                self.assertEqual(node1 << Edge(color='red', label='1.1') << node2, node2)
                self.assertEqual(node1 >> Edge(color='green', label='1.2') >> node2, node2)
                self.assertEqual(node1 << Edge(color='blue', label='1.3') >> node2, node2)

    def test_node_to_node_with_additional_attributes(self):
        with Diagram(name=os.path.join(self.name, 'node_to_node_with_additional_attributes'), show=False):
            with Cluster():
                node1 = Node("node1")
                node2 = Node("node2")
                self.assertEqual(node1 << Edge(color='red', label='2.1') << Edge(color='blue') << node2, node2)
                self.assertEqual(node1 >> Edge(color='green', label='2.2') >> Edge(color='red') >> node2, node2)
                self.assertEqual(node1 << Edge(color='blue', label='2.3') >> Edge(color='black') >> node2, node2)

    def test_nodes_to_node_with_attributes_loop(self):
        with Diagram(name=os.path.join(self.name, 'nodes_to_node_with_attributes_loop'), show=False):
            with Cluster():
                node = Node("node")
                self.assertEqual(node >> Edge(color='red', label='3.1') >> node, node)
                self.assertEqual(node << Edge(color='green', label='3.2') << node, node)
                self.assertEqual(node >> Edge(color='blue', label='3.3') << node, node)
                self.assertEqual(node << Edge(color='pink', label='3.4') >> node, node)

    def test_nodes_to_node_with_attributes_bothdirectional(self):
        with Diagram(name=os.path.join(self.name, 'nodes_to_node_with_attributes_bothdirectional'), show=False) as diagram:
            with Cluster():
                node1 = Node("node1")
                nodes = [Node("node2"), Node("node3")]
                self.assertEqual(nodes << Edge(color='green', label='4') >> node1, node1)

    def test_nodes_to_node_with_attributes_bidirectional(self):
        with Diagram(name=os.path.join(self.name, 'nodes_to_node_with_attributes_bidirectional'), show=False):
            with Cluster():
                node1 = Node("node1")
                nodes = [Node("node2"), Node("node3")]
                self.assertEqual(nodes << Edge(color='blue', label='5') >> node1, node1)

    def test_nodes_to_node_with_attributes_onedirectional(self):
        with Diagram(name=os.path.join(self.name, 'nodes_to_node_with_attributes_onedirectional'), show=False):
            with Cluster():
                node1 = Node("node1")
                nodes = [Node("node2"), Node("node3")]
                self.assertEqual(nodes >> Edge(color='red', label='6.1') >> node1, node1)
                self.assertEqual(nodes << Edge(color='green', label='6.2') << node1, node1)

    def test_nodes_to_node_with_additional_attributes_directional(self):
        with Diagram(name=os.path.join(self.name, 'nodes_to_node_with_additional_attributes_directional'), show=False):
            with Cluster():
                node1 = Node("node1")
                nodes = [Node("node2"), Node("node3")]
                self.assertEqual(nodes
                                 >> Edge(color='red', label='6.1') >> Edge(color='blue', label='6.2') >> node1, node1)
                self.assertEqual(nodes
                                 << Edge(color='green', label='6.3') << Edge(color='pink', label='6.4') << node1, node1)