update_vertex + update_edge changes

Change-Id: Icb7974c8056e1a584f6b68cb4f33f118c2e82583

vitrage/graph/driver/networkx_graph.py

@@ -13,13 +13,14 @@
 # under the License.
 
 
-from oslo_log import log as logging
-
 import copy
 import json
+
 import networkx as nx
 from networkx.readwrite import json_graph
 
+from oslo_log import log as logging
+
 from elements import Edge
 from elements import Vertex
 from graph import Direction
@@ -29,7 +30,6 @@ from vitrage.common.constants import VertexProperties as VProps
 from vitrage.graph.filter import check_filter
 from vitrage.graph.query import create_predicate
 
-
 LOG = logging.getLogger(__name__)
 
 
@@ -152,10 +152,10 @@ class NXGraph(Graph):
 
         :type v: Vertex
         """
-        if hard_update:
-            properties = self._g.node.get(v.vertex_id, None)
-            if properties:
-                properties.clear()
+        orig_prop = self._g.node.get(v.vertex_id, None)
+        if orig_prop and hard_update:
+            orig_prop.clear()
+        v.properties = self._update_properties(orig_prop, v.properties)
         self._add_vertex(v)
 
     @Notifier.update_notify
@@ -164,14 +164,27 @@ class NXGraph(Graph):
 
         :type e: Edge
         """
-        if hard_update:
-            properties = self._get_edge_properties(e.source_id,
-                                                   e.target_id,
-                                                   e.label)
-            if properties:
-                properties.clear()
+        orig_prop = self._g.edge.get(
+            e.source_id, {}).get(
+            e.target_id, {}).get(
+            e.label, None)
+        if orig_prop and hard_update:
+            orig_prop.clear()
+        e.properties = self._update_properties(orig_prop, e.properties)
         self._add_edge(e)
 
+    @staticmethod
+    def _update_properties(orig_props, new_props):
+        if orig_props is None:
+            orig_props = dict()
+        keys_to_remove = [key for key, val in new_props.items() if val is None]
+        for key in keys_to_remove:
+            del new_props[key]
+            if key in orig_props:
+                del orig_props[key]
+        orig_props.update(new_props)
+        return orig_props
+
     def remove_vertex(self, v):
         """Remove Vertex v and its edges from the graph
 

vitrage/tests/unit/graph/test_graph.py

@@ -18,7 +18,6 @@ test_vitrage graph
 
 Tests for `vitrage` graph driver
 """
-from oslo_log import log as logging
 
 from vitrage.common.constants import EdgeProperties as EProps
 from vitrage.common.constants import VertexProperties as VProps
@@ -91,6 +90,14 @@ class GraphTest(GraphTestBase):
         self.assertEqual(updated_v[VProps.CATEGORY], v[VProps.CATEGORY],
                          'Graph item should change after update')
 
+        # Update the graph item and see changes take place
+        updated_v['KUKU'] = None
+        g.update_vertex(updated_v)
+        # Get it again
+        v = g.get_vertex(v_node.vertex_id)
+        self.assertFalse('KUKU' in v.properties,
+                         'Update value to None should entirly remove the key')
+
         # check metadata
         another_vertex = utils.create_vertex(
             vitrage_id='123', entity_id='456', entity_category=INSTANCE,
@@ -170,6 +177,14 @@ class GraphTest(GraphTestBase):
                          e[EProps.UPDATE_TIMESTAMP],
                          'Graph item should change after update')
 
+        # Update the graph item and see changes take place
+        updated_e[EProps.IS_DELETED] = None
+        g.update_edge(updated_e)
+        # Get it again
+        e = g.get_edge(v_node.vertex_id, v_host.vertex_id, label)
+        self.assertFalse(EProps.IS_DELETED in e.properties,
+                         'Update value to None should entirly remove the key')
+
         # check metadata
         another_label = 'ANOTHER_LABEL'
         another_edge = utils.create_edge(