Added multi-severity to the colormap test

Author: izar

tests/dfd_colormap.dot

@@ -0,0 +1,142 @@
+digraph tm {
+    graph [
+        fontname = Arial;
+        fontsize = 14;
+    ]
+    node [
+        fontname = Arial;
+        fontsize = 14;
+        rankdir = lr;
+    ]
+    edge [
+        shape = none;
+        arrowtail = onormal;
+        fontname = Arial;
+        fontsize = 12;
+    ]
+    labelloc = "t";
+    fontsize = 20;
+    nodesep = 1;
+
+    subgraph cluster_boundary_Companynet_88f2d9c06f {
+        graph [
+            fontsize = 10;
+            fontcolor = black;
+            style = dashed;
+            color = black;
+            label = <<i>Company net</i>>;
+        ]
+
+        subgraph cluster_boundary_dmz_579e9aae81 {
+            graph [
+                fontsize = 10;
+                fontcolor = black;
+                style = dashed;
+                color = black;
+                label = <<i>dmz</i>>;
+            ]
+
+            server_Gateway_f8af758679 [
+                shape = circle;
+                color = firebrick3; fillcolor="#b2222222"; style=filled ;
+                fontcolor = black;
+                label = "Gateway";
+                margin = 0.02;
+            ]
+
+        }
+
+        subgraph cluster_boundary_backend_f2eb7a3ff7 {
+            graph [
+                fontsize = 10;
+                fontcolor = black;
+                style = dashed;
+                color = black;
+                label = <<i>backend</i>>;
+            ]
+
+            server_WebServer_2c440ebe53 [
+                shape = circle;
+                color = firebrick3; fillcolor="#b2222222"; style=filled ;
+                fontcolor = black;
+                label = "Web Server";
+                margin = 0.02;
+            ]
+
+            datastore_SQLDatabase_0291419f72 [
+                shape = none;
+                fixedsize = shape;
+                image = "/Users/izar.tarandach/Src/pytm/pytm/images/datastore_gold.png";
+                imagescale = true;
+                color = gold; fillcolor="#ffd80022"; style=filled;
+                fontcolor = black;
+                xlabel = "SQL Database";
+                label = "";
+            ]
+
+        }
+
+    }
+
+    subgraph cluster_boundary_Internet_acf3059e70 {
+        graph [
+            fontsize = 10;
+            fontcolor = black;
+            style = dashed;
+            color = black;
+            label = <<i>Internet</i>>;
+        ]
+
+        actor_User_d2006ce1bb [
+            shape = square;
+            color = darkgreen; fillcolor="#00630022"; style=filled;
+            fontcolor = black;
+            label = "User";
+            margin = 0.02;
+        ]
+
+    }
+
+    actor_User_d2006ce1bb -> server_Gateway_f8af758679 [
+        color = gold; fillcolor="#ffd80022"; style=filled;
+        fontcolor = gold; fillcolor="#ffd80022"; style=filled;
+        dir = forward;
+        label = "User enters\ncomments (*)";
+    ]
+
+    server_Gateway_f8af758679 -> server_WebServer_2c440ebe53 [
+        color = gold; fillcolor="#ffd80022"; style=filled;
+        fontcolor = gold; fillcolor="#ffd80022"; style=filled;
+        dir = forward;
+        label = "Request";
+    ]
+
+    server_WebServer_2c440ebe53 -> datastore_SQLDatabase_0291419f72 [
+        color = gold; fillcolor="#ffd80022"; style=filled;
+        fontcolor = gold; fillcolor="#ffd80022"; style=filled;
+        dir = forward;
+        label = "Insert query with\ncomments";
+    ]
+
+    datastore_SQLDatabase_0291419f72 -> server_WebServer_2c440ebe53 [
+        color = gold; fillcolor="#ffd80022"; style=filled;
+        fontcolor = gold; fillcolor="#ffd80022"; style=filled;
+        dir = forward;
+        label = "Retrieve comments";
+    ]
+
+    server_WebServer_2c440ebe53 -> server_Gateway_f8af758679 [
+        color = gold; fillcolor="#ffd80022"; style=filled;
+        fontcolor = gold; fillcolor="#ffd80022"; style=filled;
+        dir = forward;
+        label = "Response";
+    ]
+
+    server_Gateway_f8af758679 -> actor_User_d2006ce1bb [
+        color = gold; fillcolor="#ffd80022"; style=filled;
+        fontcolor = gold; fillcolor="#ffd80022"; style=filled;
+        dir = forward;
+        label = "Show comments (*)";
+    ]
+
+}

tests/test_pytmfunc.py

@@ -126,6 +126,43 @@ def test_dfd(self):
         self.maxDiff = None
         self.assertEqual(output, expected)
 
+    def test_dfd_colormap(self):
+        dir_path = os.path.dirname(os.path.realpath(__file__))
+        install_path = os.path.dirname(os.path.realpath(pytm.__file__))
+
+        with open(os.path.join(dir_path, "dfd_colormap.dot")) as x:
+            expected = (
+                x.read().strip().replace("INSTALL_PATH", os.path.dirname(install_path))
+            )
+
+        random.seed(0)
+
+        TM.reset()
+        tm = TM("my test tm", description="aaa")
+        internet = Boundary("Internet")
+        net = Boundary("Company net")
+        dmz = Boundary("dmz", inBoundary=net)
+        backend = Boundary("backend", inBoundary=net)
+        user = Actor("User", inBoundary=internet)
+        gw = Server("Gateway", inBoundary=dmz)
+        web = Server("Web Server", inBoundary=backend)
+        db = Datastore("SQL Database", inBoundary=backend, isEncryptedAtRest=True)
+        comment = Data("Comment", isStored=True)
+
+        Dataflow(user, gw, "User enters comments (*)")
+        Dataflow(gw, web, "Request")
+        Dataflow(web, db, "Insert query with comments", data=[comment])
+        Dataflow(db, web, "Retrieve comments")
+        Dataflow(web, gw, "Response")
+        Dataflow(gw, user, "Show comments (*)")
+
+        self.assertTrue(tm.check())
+        tm.resolve()
+        output = tm.dfd(colormap=True)
+
+        self.maxDiff = None
+        self.assertEqual(output, expected)
+
     def test_dfd_duplicates_ignore(self):
         dir_path = os.path.dirname(os.path.realpath(__file__))
         install_path = os.path.dirname(os.path.realpath(pytm.__file__))