Merge c7a10b742f into 1144a367b7

diagrams/__init__.py

@@ -1,8 +1,9 @@
 import contextvars
+import html
 import os
 import uuid
 from pathlib import Path
-from typing import List, Union, Dict
+from typing import List, Union, Dict, Sequence
 
 from graphviz import Digraph
 
@@ -16,10 +17,7 @@ __cluster = contextvars.ContextVar("cluster")
 
 
 def getdiagram():
-    try:
+    return __diagram.get()
-        return __diagram.get()
-    except LookupError:
-        return None
 
 
 def setdiagram(diagram):
@@ -36,9 +34,77 @@ def getcluster():
 def setcluster(cluster):
     __cluster.set(cluster)
 
+def new_init(cls, init):
+    def reset_init(*args, **kwargs):
+        cls.__init__ = init
+    return reset_init
 
-class Diagram:
+class _Cluster:
     __directions = ("TB", "BT", "LR", "RL")
+
+    def __init__(self, name=None, **kwargs):
+        self.dot = Digraph(name, **kwargs)
+        self.depth = 0
+        self.nodes = {}
+        self.subgraphs = []
+
+        try:
+            self._parent = getcluster() or getdiagram()
+        except LookupError:
+            self._parent = None
+
+    
+    def __enter__(self):
+        setcluster(self)
+        return self
+
+    def __exit__(self, *args):
+        setcluster(self._parent)
+
+        if not (self.nodes or self.subgraphs):
+            return
+
+        for node in self.nodes.values():
+            self.dot.node(node.nodeid, label=node.label, **node._attrs)
+
+        for subgraph in self.subgraphs:
+            self.dot.subgraph(subgraph.dot)
+
+        if self._parent:
+            self._parent.remove_node(self.nodeid)
+            self._parent.subgraph(self)
+
+    def node(self, node: "Node") -> None:
+        """Create a new node."""
+        self.nodes[node.nodeid] = node
+    
+    def remove_node(self, nodeid: str) -> None:
+        del self.nodes[nodeid]
+
+    def subgraph(self, subgraph: "_Cluster") -> None:
+        """Create a subgraph for clustering"""
+        self.subgraphs.append(subgraph)
+    
+    @property
+    def nodes_iter(self):
+        if self.nodes:
+            yield from self.nodes.values()
+        if self.subgraphs:
+            for subgraph in self.subgraphs:
+                yield from subgraph.nodes_iter
+
+    def _validate_direction(self, direction: str):
+        direction = direction.upper()
+        for v in self.__directions:
+            if v == direction:
+                return True
+        return False
+
+    def __str__(self) -> str:
+        return str(self.dot)
+
+
+class Diagram(_Cluster):
     __curvestyles = ("ortho", "curved")
     __outformats = ("png", "jpg", "svg", "pdf")
 
@@ -102,15 +168,19 @@ class Diagram:
         :param node_attr: Provide node_attr dot config attributes.
         :param edge_attr: Provide edge_attr dot config attributes.
         """
+
         self.name = name
         if not name and not filename:
           filename = "diagrams_image"
         elif not filename:
             filename = "_".join(self.name.split()).lower()
         self.filename = filename
-        self.dot = Digraph(self.name, filename=self.filename)
+
+        super().__init__(self.name, filename=self.filename)
+        self.edges = {}
 
         # Set attributes.
+        self.dot.attr(compound="true")
         for k, v in self._default_graph_attrs.items():
             self.dot.graph_attr[k] = v
         self.dot.graph_attr["label"] = self.name
@@ -138,29 +208,33 @@ class Diagram:
 
         self.show = show
 
-    def __str__(self) -> str:
-        return str(self.dot)
-
     def __enter__(self):
         setdiagram(self)
+        super().__enter__()
         return self
     
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(self, *args):
+        super().__exit__(*args)
+        setdiagram(None)
+
+        for (node1, node2), edge in self.edges.items():
+            cluster_node1 = next(node1.nodes_iter, None)
+            if cluster_node1:
+                edge._attrs['ltail'] = node1.nodeid
+                node1 = cluster_node1
+            cluster_node2 = next(node2.nodes_iter, None)
+            if cluster_node2:
+                edge._attrs['lhead'] = node2.nodeid
+                node2 = cluster_node2
+            self.dot.edge(node1.nodeid, node2.nodeid, **edge.attrs)
+
         self.render()
         # Remove the graphviz file leaving only the image.
         os.remove(self.filename)
-        setdiagram(None)
 
     def _repr_png_(self):
         return self.dot.pipe(format="png")
 
-    def _validate_direction(self, direction: str) -> bool:
-        direction = direction.upper()
-        for v in self.__directions:
-            if v == direction:
-                return True
-        return False
-
     def _validate_curvestyle(self, curvestyle: str) -> bool:
         curvestyle = curvestyle.lower()
         for v in self.__curvestyles:
@@ -175,24 +249,16 @@ class Diagram:
                 return True
         return False
 
-    def node(self, nodeid: str, label: str, **attrs) -> None:
-        """Create a new node."""
-        self.dot.node(nodeid, label=label, **attrs)
-
     def connect(self, node: "Node", node2: "Node", edge: "Edge") -> None:
         """Connect the two Nodes."""
-        self.dot.edge(node.nodeid, node2.nodeid, **edge.attrs)
+        self.edges[(node, node2)] = edge
-
-    def subgraph(self, dot: Digraph) -> None:
-        """Create a subgraph for clustering"""
-        self.dot.subgraph(dot)
 
     def render(self) -> None:
         self.dot.render(format=self.outformat, view=self.show, quiet=True)
 
 
-class Cluster:
+class Node(_Cluster):
-    __directions = ("TB", "BT", "LR", "RL")
+    """Node represents a node for a specific backend service."""
     __bgcolors = ("#E5F5FD", "#EBF3E7", "#ECE8F6", "#FDF7E3")
 
     # fmt: off
@@ -205,119 +271,117 @@ class Cluster:
         "fontsize": "12",
     }
 
-    # fmt: on
-
-    # FIXME:
-    #  Cluster direction does not work now. Graphviz couldn't render
-    #  correctly for a subgraph that has a different rank direction.
-    def __init__(
-        self, label: str = "cluster", direction: str = "LR", graph_attr: dict = {},
-    ):
-        """Cluster represents a cluster context.
-
-        :param label: Cluster label.
-        :param direction: Data flow direction. Default is 'left to right'.
-        :param graph_attr: Provide graph_attr dot config attributes.
-        """
-        self.label = label
-        self.name = "cluster_" + self.label
-
-        self.dot = Digraph(self.name)
-
-        # Set attributes.
-        for k, v in self._default_graph_attrs.items():
-            self.dot.graph_attr[k] = v
-        self.dot.graph_attr["label"] = self.label
-
-        if not self._validate_direction(direction):
-            raise ValueError(f'"{direction}" is not a valid direction')
-        self.dot.graph_attr["rankdir"] = direction
-
-        # Node must be belong to a diagrams.
-        self._diagram = getdiagram()
-        if self._diagram is None:
-            raise EnvironmentError("Global diagrams context not set up")
-        self._parent = getcluster()
-
-        # Set cluster depth for distinguishing the background color
-        self.depth = self._parent.depth + 1 if self._parent else 0
-        coloridx = self.depth % len(self.__bgcolors)
-        self.dot.graph_attr["bgcolor"] = self.__bgcolors[coloridx]
-
-        # Merge passed in attributes
-        self.dot.graph_attr.update(graph_attr)
-
-    def __enter__(self):
-        setcluster(self)
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        if self._parent:
-            self._parent.subgraph(self.dot)
-        else:
-            self._diagram.subgraph(self.dot)
-        setcluster(self._parent)
-
-    def _validate_direction(self, direction: str):
-        direction = direction.upper()
-        for v in self.__directions:
-            if v == direction:
-                return True
-        return False
-
-    def node(self, nodeid: str, label: str, **attrs) -> None:
-        """Create a new node in the cluster."""
-        self.dot.node(nodeid, label=label, **attrs)
-
-    def subgraph(self, dot: Digraph) -> None:
-        self.dot.subgraph(dot)
-
-
-class Node:
-    """Node represents a node for a specific backend service."""
-
     _provider = None
     _type = None
 
     _icon_dir = None
     _icon = None
-
+    _icon_size = 30
+    _direction = "LR"
     _height = 1.9
 
-    def __init__(self, label: str = "", **attrs: Dict):
+    # fmt: on
+
+    def __new__(cls, *args, **kwargs):
+        instance = object.__new__(cls)
+        lazy = kwargs.pop('_no_init', False)
+        if not lazy:
+            return instance
+        cls.__init__ = new_init(cls, cls.__init__)
+        return instance
+
+    def __init__(
+        self,
+        label: str = "",
+        direction: str = None,
+        icon: object = None,
+        icon_size: int = None,
+        **attrs: Dict
+        ):
         """Node represents a system component.
 
         :param label: Node label.
+        :param direction: Data flow direction. Default is "LR" (left to right).
+        :param icon: Custom icon for tihs cluster. Must be a node class or reference.
+        :param icon_size: The icon size when used as a Cluster. Default is 30.
         """
         # Generates an ID for identifying a node.
         self._id = self._rand_id()
-        self.label = label
+        if isinstance(label, str):
+            self.label = label
+        elif isinstance(label, Sequence):
+            self.label = "\n".join(label)
+        else:
+            self.label = str(label)
+
+        super().__init__()
+
+        if direction:
+            if not self._validate_direction(direction):
+                raise ValueError(f'"{direction}" is not a valid direction')
+            self._direction = direction
+        if icon:
+            _node = icon(_no_init=True)
+            self._icon = _node._icon
+            self._icon_dir = _node._icon_dir
+        if icon_size:
+            self._icon_size = icon_size
 
         # fmt: off
         # If a node has an icon, increase the height slightly to avoid
         # that label being spanned between icon image and white space.
         # Increase the height by the number of new lines included in the label.
-        padding = 0.4 * (label.count('\n'))
+        padding = 0.4 * (self.label.count('\n'))
+        icon_path = self._load_icon()
         self._attrs = {
             "shape": "none",
             "height": str(self._height + padding),
-            "image": self._load_icon(),
+            "image": icon_path,
-        } if self._icon else {}
+        } if icon_path else {}
+
+        self._attrs['tooltip'] = (icon if icon else self).__class__.__name__
 
         # fmt: on
         self._attrs.update(attrs)
 
-        # Node must be belong to a diagrams.
+        # If a node is in the cluster context, add it to cluster.
-        self._diagram = getdiagram()
+        if not self._parent:
-        if self._diagram is None:
             raise EnvironmentError("Global diagrams context not set up")
-        self._cluster = getcluster()
+        self._parent.node(self)
 
-        # If a node is in the cluster context, add it to cluster.
+    def __enter__(self):
-        if self._cluster:
+        super().__enter__()
-            self._cluster.node(self._id, self.label, **self._attrs)
+
+        # Set attributes.
+        for k, v in self._default_graph_attrs.items():
+            self.dot.graph_attr[k] = v
+        for k, v in self._attrs.items():
+            self.dot.graph_attr[k] = v
+
+        icon = self._load_icon()
+        if icon:
+            lines = iter(html.escape(self.label).split("\n"))
+            self.dot.graph_attr["label"] = '<<TABLE border="0"><TR>' +\
+                f'<TD fixedsize="true" width="{self._icon_size}" height="{self._icon_size}"><IMG SRC="{icon}"></IMG></TD>' +\
+                f'<TD align="left">{next(lines)}</TD></TR>' +\
+                ''.join(f'<TR><TD colspan="2" align="left">{line}</TD></TR>' for line in lines) +\
+                '</TABLE>>'
         else:
-            self._diagram.node(self._id, self.label, **self._attrs)
+            self.dot.graph_attr["label"] = self.label
+
+        self.dot.graph_attr["rankdir"] = self._direction
+
+        # Set cluster depth for distinguishing the background color
+        self.depth = self._parent.depth + 1
+        coloridx = self.depth % len(self.__bgcolors)
+        self.dot.graph_attr["bgcolor"] = self.__bgcolors[coloridx]
+
+        return self
+
+    def __exit__(self, *args):
+        super().__exit__(*args)
+        self._id = "cluster_" + self.nodeid
+        self.dot.name = self.nodeid
 
     def __repr__(self):
         _name = self.__class__.__name__
@@ -406,7 +470,7 @@ class Node:
         if not isinstance(node, Edge):
             ValueError(f"{node} is not a valid Edge")
         # An edge must be added on the global diagrams, not a cluster.
-        self._diagram.connect(self, node, edge)
+        getdiagram().connect(self, node, edge)
         return node
 
     @staticmethod
@@ -414,8 +478,10 @@ class Node:
         return uuid.uuid4().hex
 
     def _load_icon(self):
-        basedir = Path(os.path.abspath(os.path.dirname(__file__)))
+        if self._icon and self._icon_dir:
-        return os.path.join(basedir.parent, self._icon_dir, self._icon)
+            basedir = Path(os.path.abspath(os.path.dirname(__file__)))
+            return os.path.join(basedir.parent, self._icon_dir, self._icon)
+        return None
 
 
 class Edge:
@@ -464,6 +530,7 @@ class Edge:
             # Graphviz complaining about using label for edges, so replace it with xlabel.
             # Update: xlabel option causes the misaligned label position: https://github.com/mingrammer/diagrams/issues/83
             self._attrs["label"] = label
+            self._attrs["tooltip"] = label
         if color:
             self._attrs["color"] = color
         if style:
@@ -536,4 +603,4 @@ class Edge:
         return {**self._attrs, "dir": direction}
 
 
-Group = Cluster
+Group = Cluster = Node

diagrams/aws/__init__.py

@@ -2,7 +2,7 @@
 AWS provides a set of services for Amazon Web Service provider.
 """
 
-from diagrams import Node
+from diagrams import Node, Cluster
 
 
 class _AWS(Node):

diagrams/aws/cluster.py

@@ -0,0 +1,104 @@
+from diagrams import Cluster
+from diagrams.aws.compute import EC2, ApplicationAutoScaling
+from diagrams.aws.network import VPC, PrivateSubnet, PublicSubnet
+
+class Region(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "dotted",
+        "labeljust": "l",
+        "pencolor": "#AEB6BE",
+        "fontname": "Sans-Serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+
+class AvailabilityZone(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "dashed",
+        "labeljust": "l",
+        "pencolor": "#27a0ff",
+        "fontname": "sans-serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+
+class VirtualPrivateCloud(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "",
+        "labeljust": "l",
+        "pencolor": "#00D110",
+        "fontname": "sans-serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+    _icon = VPC
+
+class PrivateSubnet(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "",
+        "labeljust": "l",
+        "pencolor": "#329CFF",
+        "fontname": "sans-serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+    _icon = PrivateSubnet
+
+class PublicSubnet(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "",
+        "labeljust": "l",
+        "pencolor": "#00D110",
+        "fontname": "sans-serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+    _icon = PublicSubnet
+
+class SecurityGroup(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "dashed",
+        "labeljust": "l",
+        "pencolor": "#FF361E",
+        "fontname": "Sans-Serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+
+class AutoScalling(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "dashed",
+        "labeljust": "l",
+        "pencolor": "#FF7D1E",
+        "fontname": "Sans-Serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+    _icon = ApplicationAutoScaling
+
+class EC2Contents(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "",
+        "labeljust": "l",
+        "pencolor": "#FFB432",
+        "fontname": "Sans-Serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+    _icon = EC2

diagrams/azure/cluster.py

@@ -0,0 +1,143 @@
+from diagrams import Cluster
+from diagrams.azure.compute import VM, VMWindows, VMLinux #, VMScaleSet # Depends on PR-404
+from diagrams.azure.network import VirtualNetworks, Subnets, NetworkSecurityGroupsClassic
+
+class Subscription(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "dotted",
+        "labeljust": "l",
+        "pencolor": "#AEB6BE",
+        "fontname": "Sans-Serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+
+class Region(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "dotted",
+        "labeljust": "l",
+        "pencolor": "#AEB6BE",
+        "fontname": "Sans-Serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+
+class AvailabilityZone(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "dashed",
+        "labeljust": "l",
+        "pencolor": "#27a0ff",
+        "fontname": "sans-serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+
+class VirtualNetwork(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "",
+        "labeljust": "l",
+        "pencolor": "#00D110",
+        "fontname": "sans-serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+    _icon = VirtualNetworks
+
+class SubnetWithNSG(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "",
+        "labeljust": "l",
+        "pencolor": "#329CFF",
+        "fontname": "sans-serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+    _icon = NetworkSecurityGroupsClassic
+
+class Subnet(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "",
+        "labeljust": "l",
+        "pencolor": "#00D110",
+        "fontname": "sans-serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+    _icon = Subnets
+
+class SecurityGroup(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "dashed",
+        "labeljust": "l",
+        "pencolor": "#FF361E",
+        "fontname": "Sans-Serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+
+class VMContents(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "",
+        "labeljust": "l",
+        "pencolor": "#FFB432",
+        "fontname": "Sans-Serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+    _icon = VM
+
+class VMLinuxContents(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "",
+        "labeljust": "l",
+        "pencolor": "#FFB432",
+        "fontname": "Sans-Serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+    _icon = VMLinux
+
+class VMWindowsContents(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "",
+        "labeljust": "l",
+        "pencolor": "#FFB432",
+        "fontname": "Sans-Serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+    _icon = VMWindows
+
+# Depends on PR-404
+# class VMSS(Cluster):
+#     # fmt: off
+#     _default_graph_attrs = {
+#         "shape": "box",
+#         "style": "dashed",
+#         "labeljust": "l",
+#         "pencolor": "#FF7D1E",
+#         "fontname": "Sans-Serif",
+#         "fontsize": "12",
+#     }
+#     # fmt: on
+#     _icon = VMScaleSet

diagrams/onprem/cluster.py

@@ -0,0 +1,15 @@
+from diagrams import Cluster
+from diagrams.onprem.compute import Server
+
+class ServerContents(Cluster):
+    # fmt: off
+    _default_graph_attrs = {
+        "shape": "box",
+        "style": "rounded,dotted",
+        "labeljust": "l",
+        "pencolor": "#A0A0A0",
+        "fontname": "Sans-Serif",
+        "fontsize": "12",
+    }
+    # fmt: on
+    _icon = Server

docs/guides/cluster.md

@@ -66,6 +66,36 @@ with Diagram("Event Processing", show=False):
     handlers >> dw
 ```
 
+## Clusters with icons in the label
+
+You can add a Node icon before the cluster label (and specify its size as well).  You need to import the used Node class first.
+
+It's also possible to use the node in the `with` context adding `cluster=True` to
+make it behave like a cluster.
+
+```python
+from diagrams import Cluster, Diagram
+from diagrams.aws.compute import ECS
+from diagrams.aws.database import RDS, Aurora
+from diagrams.aws.network import Route53, VPC
+
+with Diagram("Simple Web Service with DB Cluster", show=False):
+    dns = Route53("dns")
+    web = ECS("service")
+
+    with Cluster(label='VPC',icon=VPC):
+        with Cluster("DB Cluster",icon=Aurora,icon_size=30):
+            db_master = RDS("master")
+            db_master - [RDS("slave1"),
+                         RDS("slave2")]
+        with Aurora("DB Cluster", cluster=True):
+            db_master = RDS("master")
+            db_master - [RDS("slave1"),
+                         RDS("slave2")]
+
+        dns >> web >> db_master
+```
+
 ![event processing diagram](/img/event_processing_diagram.png)
 
 > There is no depth limit of nesting. Feel free to create nested clusters as deep as you want.

examples/aws.py

@@ -0,0 +1,26 @@
+from diagrams import Diagram, Edge
+from diagrams.aws.cluster import *
+from diagrams.aws.compute import EC2
+from diagrams.onprem.container import Docker
+from diagrams.onprem.cluster import *
+from diagrams.aws.network import ELB
+
+with Diagram(name="", direction="TB", show=True):
+    with Cluster("AWS"):
+        with Region("eu-west-1"):
+            with AvailabilityZone("eu-west-1a"):
+                with VirtualPrivateCloud(""):
+                    with PrivateSubnet("Private"):
+                        with SecurityGroup("web sg"):
+                            with AutoScalling(""):
+                                with EC2Contents("A"):
+                                    d1 = Docker("Container")
+                                with ServerContents("A1"):
+                                    d2 = Docker("Container")
+
+                    with PublicSubnet("Public"):
+                        with SecurityGroup("elb sg"):
+                            lb = ELB()
+
+    lb >> Edge(forward=True, reverse=True) >> d1
+    lb >> Edge(forward=True, reverse=True) >> d2

examples/azure.py

@@ -0,0 +1,24 @@
+from diagrams import Diagram, Edge
+from diagrams.azure.cluster import *
+from diagrams.azure.compute import VM
+from diagrams.onprem.container import Docker
+from diagrams.onprem.cluster import *
+from diagrams.azure.network import LoadBalancers
+
+with Diagram(name="", filename="azure", direction="TB", show=True):
+    with Cluster("Azure"):
+        with Region("East US2"):
+            with AvailabilityZone("Zone 2"):
+                with VirtualNetwork(""):
+                    with SubnetWithNSG("Private"):
+                        # with VMScaleSet(""): # Depends on PR-404
+                        with VMContents("A"):
+                            d1 = Docker("Container")
+                        with ServerContents("A1"):
+                            d2 = Docker("Container")
+
+                    with Subnet("Public"):
+                        lb = LoadBalancers()
+
+    lb >> Edge(forward=True, reverse=True) >> d1
+    lb >> Edge(forward=True, reverse=True) >> d2

tests/test_diagram.py

@@ -135,20 +135,20 @@ class ClusterTest(unittest.TestCase):
 
     def test_with_global_context(self):
         with Diagram(name=os.path.join(self.name, "with_global_context"), show=False):
-            self.assertIsNone(getcluster())
+            self.assertEqual(getcluster(), getdiagram())
             with Cluster():
-                self.assertIsNotNone(getcluster())
+                self.assertNotEqual(getcluster(), getdiagram())
-            self.assertIsNone(getcluster())
+            self.assertEqual(getcluster(), getdiagram())
 
     def test_with_nested_cluster(self):
         with Diagram(name=os.path.join(self.name, "with_nested_cluster"), show=False):
-            self.assertIsNone(getcluster())
+            self.assertEqual(getcluster(), getdiagram())
             with Cluster() as c1:
                 self.assertEqual(c1, getcluster())
                 with Cluster() as c2:
                     self.assertEqual(c2, getcluster())
                 self.assertEqual(c1, getcluster())
-            self.assertIsNone(getcluster())
+            self.assertEqual(getcluster(), getdiagram())
 
     def test_node_not_in_diagram(self):
         # Node must be belong to a diagrams.