feat: aws nodes can be added to different drawio layers

tsypuk · Jul 28, 2023 · b58584f · b58584f
1 parent eaecf57
commit b58584f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/multicloud_diagrams/__init__.py b/multicloud_diagrams/__init__.py
@@ -34,7 +34,7 @@ class OnPrem(Enum):
 
 
 class MultiCloudDiagrams:
-    def __init__(self, debug_mode=False):
+    def __init__(self, debug_mode=False, layer_name=''):
         self.mxfile = et.Element('mxfile', host="multicloud-diagrams",
                                  agent="PIP package multicloud-diagrams. Generate resources in draw.io compatible format for Cloud infrastructure. Copyrights @ Roman Tsypuk 2023. MIT license.",
                                  type="MultiCloud")
@@ -45,12 +45,14 @@ def __init__(self, debug_mode=False):
         self.root = et.SubElement(self.mx_graph_model, 'root')
         self.layers = {}
         self.mx_cell_id_0 = et.SubElement(self.root, 'mxCell', id="0")
-        self.add_layer()
+        self.add_layer(layer_name)
         self.debug_mode = debug_mode
 
     def add_layer(self, layer_name: str = ''):
         self.layers[self.get_current_layer_count() + 1] = layer_name
         self.mx_cell_id_0 = et.SubElement(self.root, 'mxCell', id=f"{self.get_current_layer_count()}", parent="0")
+        if layer_name is not '':
+            self.mx_cell_id_0.attrib['value'] = layer_name
 
     def get_current_layer_count(self):
         return len(self.layers)
@@ -75,7 +77,7 @@ def get_layer_id(self, layer_name=None, layer_id=None):
                 raise AttributeError(f'Layer ID: {layer_id} more than current layers ({current_layers})')
             return layer_id
         else:
-            return 0
+            return 1
 
     prev_coords = {}
 
@@ -152,7 +154,7 @@ def add_service(self, id: str, node_name: str, arn: str, metadata={}, node_enum=
             raise TypeError('node_enum must be an instance of AWS,OnPrem Enum')
         self.add_vertex(id, node_name, arn, metadata, node_enum.value)
 
-    def add_vertex(self, id: str, node_name: str, arn: str, metadata: dict = {}, node_type='', layer_name='', layer_id=''):
+    def add_vertex(self, id: str, node_name: str, arn: str, metadata: dict = {}, node_type='', layer_name=None, layer_id=None):
 
         # check that there is no such vertex already
         exist = False
@@ -168,7 +170,7 @@ def add_vertex(self, id: str, node_name: str, arn: str, metadata: dict = {}, nod
 
             stringified_metadata = self.stringify_dict(metadata)
 
-            parent_id = get_layer_id(layer_name, layer_id)
+            parent_id = str(self.get_layer_id(layer_name, layer_id))
             mx_cell = et.SubElement(self.root,
                                     'mxCell',
                                     id=f'vertex:{node_type}:{id}',

diff --git a/tests/test_layers.py b/tests/test_layers.py
@@ -53,8 +53,9 @@ def verify_layers(self, mx_cells, layers_count, mcd):
         self.assertEqual(layers_count + 1, len(mx_cells))
         self.verify_mx_cell(mx_cells[0], expected={'id': '0'})
         expected_layer = {1: ''}
-        for id in range(layers_count):
-            self.verify_mx_cell(mx_cells[id + 1], expected={'id': f'{id + 1}', 'parent': '0'})
+        self.verify_mx_cell(mx_cells[1], expected={'id': f'{1}', 'parent': '0'})
+        for id in range(1, layers_count, 1):
+            self.verify_mx_cell(mx_cells[id + 1], expected={'id': f'{id + 1}', 'parent': '0', 'value': f'L{id + 1}'})
             if id > 0:
                 expected_layer[id + 1] = f'L{id + 1}'
         self.assertEqual(expected_layer, mcd.layers)