add tests for mermaid

tests/test_cli.py

@@ -122,7 +122,7 @@ def test_graphviz(load_mock, convert_mock, exit_mock):
 
 
 @patch("sys.exit")
-@patch("attack_flow.mermaid.convert")
+@patch("attack_flow.mermaid.convert_attack_flow")
 @patch("attack_flow.model.load_attack_flow_bundle")
 def test_mermaid(load_mock, convert_mock, exit_mock):
     """

tests/test_graphviz.py

@@ -6,7 +6,6 @@
     AttackCondition,
 )
 from .fixtures import get_flow_bundle, get_tree_bundle
-import json
 
 
 def test_convert_attack_flow_to_graphviz():
@@ -40,12 +39,7 @@ def test_convert_attack_flow_to_graphviz():
 
 def test_convert_attack_tree_to_graphviz():
     output = attack_flow.graphviz.convert_attack_tree(get_tree_bundle())
-    # Serializing json
-    json_object = json.dumps(output, indent=4)
-
-    # Writing to sample.json
-    with open("sample.json", "w") as outfile:
-        outfile.write(json_object)
+
     assert output == dedent(
         """\
         digraph {

tests/test_mermaid.py

@@ -1,11 +1,10 @@
 from textwrap import dedent
 
-from .fixtures import get_flow_bundle
+from .fixtures import get_flow_bundle, get_tree_bundle
 import attack_flow.mermaid
 
-
 def test_convert_attack_flow_to_mermaid():
-    output = attack_flow.mermaid.convert(get_flow_bundle())
+    output = attack_flow.mermaid.convert_attack_flow(get_flow_bundle())
     assert output == dedent(
         """\
         graph TB
@@ -41,3 +40,33 @@ class infrastructure__a75c83f7_147e_4695_b173_0981521b2f01 builtin
             attack_action__dd3820fa_bae3_4270_8000_5c4642fa780c -->|related-to| infrastructure__a75c83f7_147e_4695_b173_0981521b2f01
         """
     )
+
+def test_convert_attack_tree_to_mermaid():
+    output = attack_flow.mermaid.convert_attack_tree(get_tree_bundle())
+    assert output == dedent(
+        """\
+        graph BT
+            classDef action fill:#B40000, color:white
+            classDef AND fill:#99ccff
+            classDef OR fill:#9CE67E
+            classDef condition fill:#99ff99
+            classDef builtin fill:#cccccc
+
+            attack_action__d63857d5_1043_45a4_9397_40ef68db4c5f["<b>Action</b> - <b>Action 1</b>"]
+            class attack_action__d63857d5_1043_45a4_9397_40ef68db4c5f action
+            attack_action__1994e9f2_11f1_489a_a5e7_3ad4cfd8890a[/"<b>OR</b> - <b>My Or Operator</b>"\]
+            class attack_action__1994e9f2_11f1_489a_a5e7_3ad4cfd8890a OR
+            attack_action__24fc6003_33f6_4dd7_a929_b6031927940f["<b>Action</b> - <b>Action 2</b>"]
+            class attack_action__24fc6003_33f6_4dd7_a929_b6031927940f action
+            attack_action__a0847849_a533_4b1f_a94a_720bbd25fc17["<b>Action</b> - <b>Action 3</b>"]
+            class attack_action__a0847849_a533_4b1f_a94a_720bbd25fc17 action
+            infrastructure__79d21912_36b7_4af9_8958_38949dd0d6de["<b>Infrastructure</b> - <b>Name</b>: My<br>Infra"]
+            class infrastructure__79d21912_36b7_4af9_8958_38949dd0d6de builtin
+            attack_asset__4ae37379_6a11_44c1_b6a8_d11733cfac06["<b>Attack Asset</b> - <b>Name</b>: My<br>Asset - <b>Object Ref</b>:<br>infrastructure--<br>79d21912-36b7-4af9-8958-38949dd0d6de"]
+            class attack_asset__4ae37379_6a11_44c1_b6a8_d11733cfac06 builtin
+
+            attack_action__d63857d5_1043_45a4_9397_40ef68db4c5f -->| | attack_action__1994e9f2_11f1_489a_a5e7_3ad4cfd8890a
+            attack_action__1994e9f2_11f1_489a_a5e7_3ad4cfd8890a -->| | attack_action__a0847849_a533_4b1f_a94a_720bbd25fc17
+            attack_action__24fc6003_33f6_4dd7_a929_b6031927940f -->| | attack_action__1994e9f2_11f1_489a_a5e7_3ad4cfd8890a
+        """
+    )