feat: added fillColor change for any AWS node; added tests

multicloud_diagrams/__init__.py

@@ -5,6 +5,7 @@
 import pkgutil
 from enum import Enum
 from typing import Union
+import re
 
 import yaml
 
@@ -33,9 +34,24 @@ class OnPrem(Enum):
 Services = Union[AWS, OnPrem]
 
 
+def update_fill_color(style_str, node_color):
+    pattern = r'fillColor=([^;]+)'
+    # Replace the 'fillColor' value with the node_color
+    return re.sub(pattern, f'fillColor={node_color}', style_str)
+
+
+def stringify_dict(metadata: dict) -> str:
+    if metadata:
+        # return '<BR>-----------<BR>' + '<BR>'.join([f'<b>{k.capitalize()}</b>: {v}' for k, v in metadata.items()])
+        return '<BR>-----------<BR>' + '<BR>'.join([f'<b>{k}</b>: {v}' for k, v in metadata.items()])
+    else:
+        return ''
+
+
 class MultiCloudDiagrams:
     def __init__(self, debug_mode=False, shadow=True, layer_name=''):
-        self.mxfile = et.Element('mxfile', host="multicloud-diagrams",
+        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")
 
@@ -97,19 +113,11 @@ def get_layer_id(self, layer_name=None, layer_id=None):
 
     def get_vertex_metadata(self, node_type: str) -> dict:
         if node_type in self.supported_vertex:
-            return self.supported_vertex[node_type]
+            return self.supported_vertex[node_type].copy()
         else:
             logging.warning(
                 f'No such nodeType: {node_type} in the Library (using default fallback icon Info). Please contact maintainer to add it, or provide MergeRequest')
-            return self.supported_vertex['fallback_vertex']
-
-    @staticmethod
-    def stringify_dict(metadata: dict) -> str:
-        if metadata:
-            # return '<BR>-----------<BR>' + '<BR>'.join([f'<b>{k.capitalize()}</b>: {v}' for k, v in metadata.items()])
-            return '<BR>-----------<BR>' + '<BR>'.join([f'<b>{k}</b>: {v}' for k, v in metadata.items()])
-        else:
-            return ''
+            return self.supported_vertex['fallback_vertex'].copy()
 
     # green fill_color=#d5e8d4
     # red fill_color=#f8cecc;"
@@ -160,7 +168,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=None, layer_id=None):
+    def add_vertex(self, id: str, node_name: str, arn: str, metadata: dict = {}, node_type='', layer_name=None, layer_id=None, fill_color=None):
 
         # check that there is no such vertex already
         exist = False
@@ -173,8 +181,10 @@ def add_vertex(self, id: str, node_name: str, arn: str, metadata: dict = {}, nod
 
         if not exist:
             shape_parameters = self.get_vertex_metadata(node_type)
+            if fill_color is not None:
+                shape_parameters['style'] = update_fill_color(shape_parameters['style'], fill_color)
 
-            stringified_metadata = self.stringify_dict(metadata)
+            stringified_metadata = stringify_dict(metadata)
 
             parent_id = str(self.get_layer_id(layer_name, layer_id))
             mx_cell = et.SubElement(self.root,

tests/test___init__.py

@@ -1,6 +1,6 @@
 from unittest import TestCase
 
-from multicloud_diagrams import MultiCloudDiagrams
+from multicloud_diagrams import MultiCloudDiagrams, update_fill_color, stringify_dict
 
 
 class TestMultiCloudDiagrams(TestCase):
@@ -70,18 +70,26 @@ def test__build_vertex_id_approach_3_dst(self):
 
     def test_stringify_dict(self):
         # given
-        mcd = MultiCloudDiagrams()
         metadata = {'key1': 'value1', 'key2': 'value2', 'key3': 'value3'}
         # when
-        result = mcd.stringify_dict(metadata)
+        result = stringify_dict(metadata)
         # then
         self.assertEqual('<BR>-----------<BR><b>key1</b>: value1<BR><b>key2</b>: value2<BR><b>key3</b>: value3', result)
 
     def test_stringify_dict_empty(self):
         # given
-        mcd = MultiCloudDiagrams()
         metadata = {}
         # when
-        result = mcd.stringify_dict(metadata)
+        result = stringify_dict(metadata)
         # then
         self.assertEqual('', result)
+
+    def test_update_fill_color(self):
+        # given-when
+        result = update_fill_color(
+            style_str='sketch=0;outlineConnect=0;fontColor=#232F3E;gradientColor=#FF4F8B;gradientDirection=north;fillColor=#BC1356;strokeColor=#ffffff;dashed=0;verticalLabelPosition=bottom;verticalAlign=top;align=left;html=1;fontSize=12;fontStyle=0;aspect=fixed;shape=mxgraph.aws4.resourceIcon;resIcon=mxgraph.aws4.sns;',
+            node_color='#FF0000')
+        # then
+        self.assertEqual(
+            'sketch=0;outlineConnect=0;fontColor=#232F3E;gradientColor=#FF4F8B;gradientDirection=north;fillColor=#FF0000;strokeColor=#ffffff;dashed=0;verticalLabelPosition=bottom;verticalAlign=top;align=left;html=1;fontSize=12;fontStyle=0;aspect=fixed;shape=mxgraph.aws4.resourceIcon;resIcon=mxgraph.aws4.sns;',
+            result)

tests/test_colors.py

@@ -0,0 +1,52 @@
+from multicloud_diagrams import MultiCloudDiagrams
+
+from utils.utils import TestUtilities
+
+
+class TestMultiCloudDiagramsColors(TestUtilities):
+
+    def test_add_default_node_color(self):
+        # given
+        mcd = MultiCloudDiagrams()
+
+        # when
+        sns_arn = 'arn:aws:sns:eu-west-1:123456789012:internal.fifo'
+        metadata = {
+            "Owner": 123456789012,
+            "SubscriptionsConfirmed": 3,
+            "SubscriptionsPending": 0
+        }
+        mcd.add_vertex(id=sns_arn, node_name='internal.fifo', arn=sns_arn, node_type='sns', metadata=metadata)
+
+        # then
+        expected = {
+            'id': 'vertex:sns:arn:aws:sns:eu-west-1:123456789012:internal.fifo',
+            'value': '<b>Name</b>: internal.fifo<BR><b>ARN</b>: arn:aws:sns:eu-west-1:123456789012:internal.fifo <BR>-----------<BR><b>Owner</b>: 123456789012<BR><b>SubscriptionsConfirmed</b>: 3<BR><b>SubscriptionsPending</b>: 0',
+            'parent': '1',
+            'vertex': '1'
+        }
+        self.verify_aws_resource(expected, mcd.mxfile, 'internal.fifo', 'sns')
+
+    def test_add_custom_node_color(self):
+        # given
+        mcd = MultiCloudDiagrams()
+        node_color = '#FF0000'
+
+        sns_arn = 'arn:aws:sns:eu-west-1:123456789012:internal.fifo'
+        metadata = {
+            "Owner": 123456789012,
+            "SubscriptionsConfirmed": 3,
+            "SubscriptionsPending": 0
+        }
+
+        # when
+        mcd.add_vertex(id=sns_arn, node_name='internal.fifo', arn=sns_arn, node_type='sns', metadata=metadata, fill_color=node_color)
+
+        # then
+        expected = {
+            'id': 'vertex:sns:arn:aws:sns:eu-west-1:123456789012:internal.fifo',
+            'value': '<b>Name</b>: internal.fifo<BR><b>ARN</b>: arn:aws:sns:eu-west-1:123456789012:internal.fifo <BR>-----------<BR><b>Owner</b>: 123456789012<BR><b>SubscriptionsConfirmed</b>: 3<BR><b>SubscriptionsPending</b>: 0',
+            'parent': '1',
+            'vertex': '1'
+        }
+        self.verify_aws_resource(expected, mcd.mxfile, 'internal.fifo', 'sns', fill_color=node_color)

tests/test_layers.py

@@ -49,17 +49,6 @@ def test_add_3_layers(self):
 
         self.verify_layers(mx_cells, 3, mcd)
 
-    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: ''}
-        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)
-
     def test_get_layer_id_single_layer(self):
         # given
         mcd = MultiCloudDiagrams(debug_mode=True)

utils/utils.py

@@ -3,6 +3,8 @@
 import json
 import xml.etree.ElementTree as et
 
+from multicloud_diagrams import update_fill_color
+
 
 class TestUtilities(unittest.TestCase):
     supported_vertex = {}
@@ -74,7 +76,7 @@ def verify_vertex_in_isolation(self, mx_cells):
         self.verify_mx_cell(mx_cells[0], expected={'id': '0'})
         self.verify_mx_cell(mx_cells[1], expected={'id': '1', 'parent': '0'})
 
-    def verify_aws_resource(self, expected: dict, mx_file: et.Element, resource_name, resource_type, debug_mode=False):
+    def verify_aws_resource(self, expected: dict, mx_file: et.Element, resource_name, resource_type, debug_mode=False, fill_color=None):
         tree = et.ElementTree(mx_file)
         self.verify_mxfile_default(et.ElementTree(tree))
 
@@ -85,6 +87,8 @@ def verify_aws_resource(self, expected: dict, mx_file: et.Element, resource_name
         if 'style' in expected:
             del expected['style']
         expected['style'] = self.supported_vertex[resource_type]['style']
+        if fill_color is not None:
+            expected['style'] = update_fill_color(expected['style'], fill_color)
 
         children = mx_cells[2].findall("./*")
 
@@ -105,3 +109,14 @@ def verify_aws_resource(self, expected: dict, mx_file: et.Element, resource_name
         # 'height', 'width' are verified based on providers file content
         self.assertEqual(self.supported_vertex[resource_type]['height'], mx_geometry.attrib['height'])
         self.assertEqual(self.supported_vertex[resource_type]['width'], mx_geometry.attrib['width'])
+
+    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: ''}
+        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)