diff --git a/tests/test_diagram.py b/tests/test_diagram.py
index bbb1aa5..ad8558c 100644
--- a/tests/test_diagram.py
+++ b/tests/test_diagram.py
@@ -2,7 +2,7 @@ import os
 import shutil
 import unittest
 
-from diagrams import Cluster, Diagram, Node, Edge
+from diagrams import Cluster, Diagram, Edge, Node
 from diagrams import getcluster, getdiagram, setcluster, setdiagram
 
 
@@ -101,11 +101,11 @@ class DiagramTest(unittest.TestCase):
         self.assertTrue(os.path.exists(f"{self.name}.png"))
 
     def test_empty_name(self):
-      """Check that providing an empty name don't crash, but save in a diagrams_image.xxx file."""
-      self.name = 'diagrams_image'
-      with Diagram(show=False):
-          Node("node1")
-      self.assertTrue(os.path.exists(f"{self.name}.png"))
+        """Check that providing an empty name don't crash, but save in a diagrams_image.xxx file."""
+        self.name = 'diagrams_image'
+        with Diagram(show=False):
+            Node("node1")
+        self.assertTrue(os.path.exists(f"{self.name}.png"))
 
 
 class ClusterTest(unittest.TestCase):
@@ -251,9 +251,7 @@ class EdgeTest(unittest.TestCase):
                 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 Diagram(name=os.path.join(self.name, "nodes_to_node_with_attributes_bothdirectional"), show=False):
             with Cluster():
                 node1 = Node("node1")
                 nodes = [Node("node2"), Node("node3")]