Allow interchangeable graph class for the entity graph.

Graph driver class is a configurable entry point using stevedore, can be changed in entry_points.txt and vitrage.conf.
NetworkX graph - classes encapsulation, can now easily be replaced

Change-Id: I7390d2f692886c0e2c67159b63486e2956046bfd

setup.cfg

@@ -29,6 +29,9 @@ console_scripts =
     vitrage-graph = vitrage.cmd.graph:main
     vitrage-notifier = vitrage.cmd.notifier:main
 
+vitrage.entity_graph =
+    networkx = vitrage.graph.driver.networkx_graph:NXGraph
+
 oslo.config.opts =
     vitrage = vitrage.opts:list_opts
     datasources = vitrage.opts:datasources_opts

vitrage/api_handler/apis/rca.py

@@ -18,7 +18,6 @@ from vitrage.api_handler.apis.base import ALARMS_ALL_QUERY
 from vitrage.api_handler.apis.base import EDGE_QUERY
 from vitrage.api_handler.apis.base import EntityGraphApisBase
 from vitrage.api_handler.apis.base import RCA_QUERY
-from vitrage.graph import create_algorithm
 from vitrage.graph import Direction
 
 
@@ -37,7 +36,7 @@ class RcaApis(EntityGraphApisBase):
 
         project_id = ctx.get(self.TENANT_PROPERTY, None)
         is_admin_project = ctx.get(self.IS_ADMIN_PROJECT_PROPERTY, False)
-        ga = create_algorithm(self.entity_graph)
+        ga = self.entity_graph.algo
 
         found_graph_out = ga.graph_query_vertices(query_dict=RCA_QUERY,
                                                   root_id=root,

vitrage/api_handler/apis/topology.py

@@ -24,7 +24,6 @@ from vitrage.datasources.nova.instance import NOVA_INSTANCE_DATASOURCE
 from vitrage.datasources import OPENSTACK_CLUSTER
 from vitrage.datasources.transformer_base import build_key
 from vitrage.datasources.transformer_base import CLUSTER_ID
-from vitrage.graph import create_algorithm
 
 
 LOG = log.getLogger(__name__)
@@ -42,7 +41,7 @@ class TopologyApis(EntityGraphApisBase):
 
         project_id = ctx.get(self.TENANT_PROPERTY, None)
         is_admin_project = ctx.get(self.IS_ADMIN_PROJECT_PROPERTY, False)
-        ga = create_algorithm(self.entity_graph)
+        ga = self.entity_graph.algo
 
         if graph_type == 'tree':
             if not query:

vitrage/cmd/graph.py

@@ -24,12 +24,12 @@ from vitrage.common.constants import EntityCategory
 from vitrage.datasources import launcher as datasource_launcher
 from vitrage.datasources import OPENSTACK_CLUSTER
 from vitrage.datasources.transformer_base import CLUSTER_ID
+from vitrage import entity_graph
 from vitrage.entity_graph.consistency import service as consistency_svc
 from vitrage.entity_graph.initialization_status import InitializationStatus
 from vitrage.entity_graph import service as entity_graph_svc
 from vitrage.evaluator.scenario_evaluator import ScenarioEvaluator
 from vitrage.evaluator.scenario_repository import ScenarioRepository
-from vitrage.graph import create_graph
 from vitrage import service
 
 
@@ -67,7 +67,7 @@ def main():
 def init(conf):
     mp_queue = multiprocessing.Queue()
     evaluator_q = queue.Queue()
-    e_graph = create_graph(
+    e_graph = entity_graph.get_graph_driver(conf)(
         'Entity Graph',
         '%s:%s:%s' % (EntityCategory.RESOURCE, OPENSTACK_CLUSTER, CLUSTER_ID))
     scenario_repo = ScenarioRepository(conf)

vitrage/common/utils.py

@@ -19,6 +19,8 @@
 
 from oslo_config import cfg
 
+import cProfile
+
 
 def recursive_keypairs(d, separator='.'):
     # taken from ceilometer and gnocchi
@@ -35,3 +37,16 @@ def opt_exists(conf_parent, opt):
         return conf_parent[opt]
     except cfg.NoSuchOptError:
         return False
+
+
+def do_cprofile(func):
+    def profiled_func(*args, **kwargs):
+        profile = cProfile.Profile()
+        try:
+            profile.enable()
+            result = func(*args, **kwargs)
+            profile.disable()
+            return result
+        finally:
+            profile.print_stats('cumulative')
+    return profiled_func

vitrage/datasources/launcher.py

@@ -41,6 +41,7 @@ class Launcher(object):
         self.services = self._register_services()
 
     def launch(self):
+        # launcher = os_service.ServiceLauncher(self.conf)  # For Debugging
         launcher = os_service.ProcessLauncher(self.conf)
         for service in self.services:
             launcher.launch_service(service, 1)

vitrage/entity_graph/__init__.py

@@ -15,6 +15,7 @@
 
 from oslo_config import cfg
 
+from stevedore import driver
 
 OPTS = [
     cfg.StrOpt('datasources_values_dir',
@@ -26,4 +27,17 @@ OPTS = [
                default='vitrage.graph',
                help='The topic that vitrage-graph uses for graph '
                     'notification messages.'),
+    cfg.StrOpt('graph_driver',
+               default='networkx',
+               help='graph driver implementation class'),
 ]
+
+
+def get_graph_driver(conf):
+    try:
+        mgr = driver.DriverManager('vitrage.entity_graph',
+                                   conf.entity_graph.graph_driver,
+                                   invoke_on_load=True)
+        return mgr.driver
+    except ImportError:
+        return None

vitrage/entity_graph/processor/processor.py

@@ -24,8 +24,8 @@ from vitrage.entity_graph.processor import base as processor
 from vitrage.entity_graph.processor.notifier import GraphNotifier
 from vitrage.entity_graph.processor import processor_utils as PUtils
 from vitrage.entity_graph.transformer_manager import TransformerManager
-from vitrage.graph import create_graph
 from vitrage.graph import Direction
+from vitrage.graph.driver.networkx_graph import NXGraph
 
 LOG = log.getLogger(__name__)
 
@@ -40,7 +40,7 @@ class Processor(processor.ProcessorBase):
         self._initialize_events_actions()
         self.initialization_status = initialization_status
         self.entity_graph = e_graph if e_graph is not None\
-            else create_graph("Entity Graph")
+            else NXGraph("Entity Graph")
         self._notifier = GraphNotifier(conf)
 
     def process_event(self, event):

vitrage/evaluator/scenario_evaluator.py

@@ -28,8 +28,7 @@ from vitrage.evaluator.template_data import ActionSpecs
 from vitrage.evaluator.template_data import EdgeDescription
 from vitrage.evaluator.template_data import ENTITY
 from vitrage.graph.algo_driver.algorithm import Mapping
-from vitrage.graph import create_algorithm
+from vitrage.graph.driver.networkx_graph import NXGraph
-from vitrage.graph import create_graph
 from vitrage.graph.driver import Vertex
 
 LOG = log.getLogger(__name__)
@@ -55,7 +54,6 @@ class ScenarioEvaluator(object):
         self.conf = conf
         self._scenario_repo = scenario_repo
         self._entity_graph = entity_graph
-        self._graph_algs = create_algorithm(entity_graph)
         self._action_executor = ActionExecutor(event_queue)
         self._entity_graph.subscribe(self.process_event)
         self._action_tracker = ActionTracker(DatasourceInfoMapper(self.conf))
@@ -204,7 +202,7 @@ class ScenarioEvaluator(object):
 
     def _evaluate_and_condition(self, condition, element, scenario_element):
 
-        condition_g = create_graph("scenario condition")
+        condition_g = NXGraph("scenario condition")
         for term in condition:
             if not term.positive:
                 # todo(erosensw): add support for NOT clauses
@@ -224,7 +222,8 @@ class ScenarioEvaluator(object):
             initial_map = Mapping(scenario_element, element, True)
         else:
             initial_map = Mapping(scenario_element.edge, element, False)
-        return self._graph_algs.sub_graph_matching(condition_g, [initial_map])
+        return self._entity_graph.algo.sub_graph_matching(condition_g,
+                                                          [initial_map])
 
     @staticmethod
     def _set_relationship_not_deleted(edge_description):

vitrage/graph/algo_driver/__init__.py

@@ -11,18 +11,4 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 # License for the specific language governing permissions and limitations
 # under the License.
-
+__author__ = 'stack'
-from vitrage.graph.algo_driver.algorithm import *  # noqa
-from vitrage.graph.algo_driver.networkx_algorithm import NXAlgorithm
-
-
-def create_algorithm(graph):
-    """Create a Graph algorithm instance
-
-    For now only return NXAlgorithm
-
-    :param graph:
-    :type graph: Graph
-    :rtype: GraphAlgorithm
-    """
-    return NXAlgorithm(graph=graph)

vitrage/graph/algo_driver/networkx_algorithm.py

@@ -21,7 +21,6 @@ from vitrage.graph.algo_driver.algorithm import GraphAlgorithm
 from vitrage.graph.algo_driver.sub_graph_matching import subgraph_matching
 from vitrage.graph.driver import Direction
 from vitrage.graph.driver import Edge
-from vitrage.graph.driver import NXGraph
 from vitrage.graph.driver import Vertex
 from vitrage.graph.filter import check_filter
 from vitrage.graph.query import create_predicate
@@ -39,14 +38,18 @@ class NXAlgorithm(GraphAlgorithm):
         """
         super(NXAlgorithm, self).__init__(graph)
 
+    @classmethod
+    def _create_new_graph(cls, *args, **kwargs):
+        from vitrage.graph.driver.networkx_graph import NXGraph
+        return NXGraph(args, **kwargs)
+
     def graph_query_vertices(self,
                              query_dict=None,
                              root_id=None,
                              depth=None,
                              direction=Direction.BOTH,
                              edge_query_dict=None):
-
+        graph = self._create_new_graph('graph')
-        graph = NXGraph('graph')
 
         if not root_id:
             root_id = self.graph.root_id
@@ -80,7 +83,8 @@ class NXAlgorithm(GraphAlgorithm):
             e_result.extend(e_list)
             nodes_q.extend([(v_id, curr_depth + 1) for v_id, data in n_list])
 
-        graph = NXGraph(graph.name,
+        graph = self._create_new_graph(
+            graph.name,
             graph.root_id,
             vertices=self._vertex_result_to_list(n_result),
             edges=self._edge_result_to_list(e_result))
@@ -125,7 +129,7 @@ class NXAlgorithm(GraphAlgorithm):
 
         vertices_ids = [vertex.vertex_id for vertex in vertices]
 
-        graph = NXGraph('graph')
+        graph = self._create_new_graph('graph')
         graph._g = self.graph._g.subgraph(vertices_ids)
 
         # delete non matching edges
@@ -143,7 +147,7 @@ class NXAlgorithm(GraphAlgorithm):
         return graph
 
     def subgraph(self, entities):
-        subgraph = NXGraph('graph')
+        subgraph = self._create_new_graph('graph')
         subgraph._g = self.graph._g.subgraph(entities)
         return subgraph
 

vitrage/graph/driver/__init__.py

@@ -14,16 +14,3 @@
 
 from vitrage.graph.driver.elements import *  # noqa
 from vitrage.graph.driver.graph import *  # noqa
-from vitrage.graph.driver.networkx_graph import NXGraph
-
-
-def create_graph(name, root_id=None):
-    """Create a Graph instance
-
-    For now only return NXGraph
-
-    :param root_id:
-    :type name: str
-    :rtype: Graph
-    """
-    return NXGraph(name, root_id)

vitrage/graph/driver/graph.py

@@ -20,6 +20,7 @@ vitrage.graph.driver namespace.
 
 """
 import abc
+import copy
 import six
 
 from vitrage.graph.driver.elements import Edge
@@ -61,6 +62,14 @@ class Graph(object):
         if isinstance(item, Vertex):
             return self.get_vertex(item.vertex_id)
 
+    @property
+    def algo(self):
+        """Get graph algorithms
+
+        :rtype: GraphAlgorithm
+        """
+        return None
+
     @abc.abstractmethod
     def copy(self):
         """Create a copy of the graph
@@ -105,7 +114,6 @@ class Graph(object):
         """
         pass
 
-    @abc.abstractmethod
     def add_vertices(self, vertices):
         """Add a list of vertices to the graph
 
@@ -114,7 +122,11 @@ class Graph(object):
         :param vertices:
         :type vertices:list of Vertex
         """
-        pass
+        if not vertices:
+            return
+
+        for v in vertices:
+            self.add_vertex(v)
 
     @abc.abstractmethod
     def add_edge(self, e):
@@ -143,7 +155,6 @@ class Graph(object):
         """
         pass
 
-    @abc.abstractmethod
     def add_edges(self, edges):
         """Add a list of edges to the graph
 
@@ -152,7 +163,11 @@ class Graph(object):
         :param edges:
         :type edges:list of Edge
         """
-        pass
+        if not edges:
+            return
+
+        for e in edges:
+            self.add_edge(e)
 
     @abc.abstractmethod
     def get_vertex(self, v_id):
@@ -213,53 +228,52 @@ class Graph(object):
         :type attr_filter: dict
 
         :return: All edges matching the requirements
-        :rtype: list of Edge
+        :rtype: set of Edge
         """
         pass
 
     @abc.abstractmethod
-    def update_vertex(self, v, hard_update=False):
+    def update_vertex(self, v):
         """Update the vertex properties
 
         Update an existing vertex and create it if non existing.
-        Hard update: can be used to remove existing fields.
 
         :param v: the vertex with the new data
         :type v: Vertex
-        :param hard_update: if True, original properties will be removed.
-        :type hard_update: bool
         """
         pass
 
-    @abc.abstractmethod
+    def update_vertices(self, vertices):
-    def update_vertices(self, vertices, hard_update=False):
         """For each vertex, update its properties
 
         For each existing vertex, update its properties and create it if
         non existing.
-        Hard update: can be used to remove existing fields.
 
         :param vertices: the vertex with the new data
         :type vertices: List
-        :param hard_update: if True, original properties will be removed.
-        :type hard_update: bool
         """
-        pass
+        for v in vertices:
+            self.update_vertex(v)
 
     @abc.abstractmethod
-    def update_edge(self, e, hard_update=False):
+    def update_edge(self, e):
         """Update the edge properties
 
         Update an existing edge and create it if non existing.
-        Hard update: can be used to remove existing fields.
 
         :param e: the edge with the new data
         :type e: Edge
-        :param hard_update: if True, original properties will be removed.
-        :type hard_update: bool
         """
         pass
 
+    @staticmethod
+    def _merge_properties(base_props, new_props):
+        if base_props is None:
+            base_props = copy.copy(new_props)
+        else:
+            base_props.update(copy.copy(new_props))
+        return {k: v for k, v in base_props.items() if v is not None}
+
     @abc.abstractmethod
     def remove_vertex(self, v):
         """Remove Vertex v and its edges from the graph

vitrage/graph/driver/networkx_graph.py

@@ -21,6 +21,7 @@ from networkx.readwrite import json_graph
 from oslo_log import log as logging
 
 from vitrage.common.constants import VertexProperties as VProps
+from vitrage.graph.algo_driver.networkx_algorithm import NXAlgorithm
 from vitrage.graph.driver.elements import Edge
 from vitrage.graph.driver.elements import Vertex
 from vitrage.graph.driver.graph import Direction
@@ -59,6 +60,10 @@ class NXGraph(Graph):
     def __len__(self):
         return len(self._g)
 
+    @property
+    def algo(self):
+        return NXAlgorithm(self)
+
     def copy(self):
         self_copy = NXGraph(self.name, self.root_id)
         self_copy._g = self._g.copy()
@@ -77,13 +82,6 @@ class NXGraph(Graph):
         properties_copy = copy.copy(v.properties)
         self._g.add_node(n=v.vertex_id, attr_dict=properties_copy)
 
-    def add_vertices(self, vertices):
-        if not vertices:
-            return
-
-        for v in vertices:
-            self.add_vertex(v)
-
     @Notifier.update_notify
     def add_edge(self, e):
         """Add an edge to the graph
@@ -98,13 +96,6 @@ class NXGraph(Graph):
         self._g.add_edge(u=e.source_id, v=e.target_id,
                          key=e.label, attr_dict=properties_copy)
 
-    def add_edges(self, edges):
-        if not edges:
-            return
-
-        for e in edges:
-            self.add_edge(e)
-
     def get_vertex(self, v_id):
         """Fetch a vertex from the graph
 
@@ -131,15 +122,15 @@ class NXGraph(Graph):
                   attr_filter=None):
         """Fetch multiple edges from the graph
 
-        :rtype: list of Edge
+        :rtype: set of Edge
         """
         def check_edge(edge_data):
             return check_filter(edge_data, attr_filter)
 
         nodes, edges = self._neighboring_nodes_edges_query(
             v_id, edge_predicate=check_edge, direction=direction)
-        edge_copies = [edge_copy(u, v, label, data)
+        edge_copies = set(edge_copy(u, v, label, data)
-                       for u, v, label, data in edges]
+                          for u, v, label, data in edges)
         return edge_copies
 
     def _get_edges_by_direction(self, v_id, direction):
@@ -165,33 +156,22 @@ class NXGraph(Graph):
         return self._g.number_of_edges()
 
     @Notifier.update_notify
-    def update_vertex(self, v, hard_update=False):
+    def update_vertex(self, v):
         """Update the vertex properties
 
-        :param hard_update:
         :type v: Vertex
         """
         orig_prop = self._g.node.get(v.vertex_id, None)
         if not orig_prop:
             self._add_vertex(v)
             return
-        new_prop = self._merge_properties(orig_prop, v.properties, hard_update)
+        new_prop = self._merge_properties(orig_prop, v.properties)
         self._g.node[v.vertex_id] = new_prop
 
-    def update_vertices(self, vertices, hard_update=False):
-        """For each vertex, update its properties
-
-        :param hard_update:
-        :type vertices: List
-        """
-        for v in vertices:
-            self.update_vertex(v, hard_update)
-
     @Notifier.update_notify
-    def update_edge(self, e, hard_update=False):
+    def update_edge(self, e):
         """Update the edge properties
 
-        :param hard_update:
         :type e: Edge
         """
         orig_prop = self._g.edge.get(
@@ -201,17 +181,9 @@ class NXGraph(Graph):
         if not orig_prop:
             self._add_edge(e)
             return
-        new_prop = self._merge_properties(orig_prop, e.properties, hard_update)
+        new_prop = self._merge_properties(orig_prop, e.properties)
         self._g.edge[e.source_id][e.target_id][e.label] = new_prop
 
-    @staticmethod
-    def _merge_properties(base_props, new_props, hard_update):
-        if base_props is None or hard_update:
-            base_props = copy.copy(new_props)
-        else:
-            base_props.update(copy.copy(new_props))
-        return {k: v for k, v in base_props.items() if v is not None}
-
     def remove_vertex(self, v):
         """Remove Vertex v and its edges from the graph
 

vitrage/graph/driver/notifier.py

@@ -27,8 +27,8 @@ def _after_func(graph, item, data_before=None):
     if not graph.is_subscribed():
         return
     element = graph.get_item(item)
-    is_vertex = isinstance(element, Vertex)
+    is_vertex = isinstance(element, Vertex) or isinstance(item, Vertex)
-    graph.notifier.notify(data_before, graph.get_item(item), is_vertex, graph)
+    graph.notifier.notify(data_before, element, is_vertex, graph)
 
 
 class Notifier(object):

vitrage/tests/functional/api_handler/test_apis.py

@@ -23,7 +23,7 @@ from vitrage.datasources import NOVA_HOST_DATASOURCE
 from vitrage.datasources import NOVA_INSTANCE_DATASOURCE
 from vitrage.datasources import NOVA_ZONE_DATASOURCE
 from vitrage.datasources import OPENSTACK_CLUSTER
-from vitrage.graph import NXGraph
+from vitrage.graph.driver.networkx_graph import NXGraph
 import vitrage.graph.utils as graph_utils
 from vitrage.tests.unit.entity_graph.base import TestEntityGraphUnitBase
 

vitrage/tests/unit/entity_graph/processor/test_entity_graph.py

@@ -13,7 +13,7 @@
 # under the License.
 
 from vitrage.entity_graph.processor import processor_utils as PUtils
-from vitrage.graph import NXGraph
+from vitrage.graph.driver.networkx_graph import NXGraph
 from vitrage.tests.unit.entity_graph.processor import base
 
 

vitrage/tests/unit/graph/base.py

@@ -26,12 +26,13 @@ from oslo_log import log as logging
 
 from vitrage.common.constants import EdgeLabel as ELabel
 from vitrage.common.constants import EntityCategory
+from vitrage.common.exception import VitrageError
 from vitrage.datasources.nova.host import NOVA_HOST_DATASOURCE
 from vitrage.datasources.nova.instance import NOVA_INSTANCE_DATASOURCE
 from vitrage.datasources import OPENSTACK_CLUSTER
 from vitrage.datasources.static_physical import SWITCH
 from vitrage.datasources.transformer_base import CLUSTER_ID
-from vitrage.graph import create_graph
+from vitrage.graph.driver.networkx_graph import NXGraph
 from vitrage.graph import utils as graph_utils
 from vitrage.tests import base
 
@@ -120,31 +121,19 @@ class GraphTestBase(base.BaseTest):
     def __init__(self, *args, **kwds):
         super(GraphTestBase, self).__init__(*args, **kwds)
 
-        self.vm_id = 10000000
-        self.vm_alarm_id = 30000000
-        self.vms = []
-        self.host_alarm_id = 20000000
-        self.host_test_id = 40000000
-        self.entity_graph = self._create_entity_graph(
-            'entity_graph',
-            num_of_hosts_per_node=ENTITY_GRAPH_HOSTS_PER_CLUSTER,
-            num_of_vms_per_host=ENTITY_GRAPH_VMS_PER_HOST,
-            num_of_alarms_per_host=ENTITY_GRAPH_ALARMS_PER_HOST,
-            num_of_alarms_per_vm=ENTITY_GRAPH_ALARMS_PER_VM,
-            num_of_tests_per_host=ENTITY_GRAPH_TESTS_PER_HOST)
-
     def _assert_set_equal(self, d1, d2, message):
         super(GraphTestBase, self).assert_dict_equal(
             dict.fromkeys(d1, 0), dict.fromkeys(d2, 0), message)
 
-    def _create_entity_graph(self, name, num_of_alarms_per_host,
+    @classmethod
+    def _create_entity_graph(cls, name, num_of_alarms_per_host,
                              num_of_alarms_per_vm,
                              num_of_hosts_per_node,
                              num_of_vms_per_host,
                              num_of_tests_per_host):
 
         start = time.time()
-        g = create_graph(name, EntityCategory.RESOURCE + ':' +
+        g = NXGraph(name, EntityCategory.RESOURCE + ':' +
                     OPENSTACK_CLUSTER + ':' +
                     CLUSTER_ID)
         g.add_vertex(v_node)
@@ -167,34 +156,34 @@ class GraphTestBase(base.BaseTest):
             # Add Host Alarms
             for j in range(num_of_alarms_per_host):
                 add_connected_vertex(g, ALARM, ALARM_ON_HOST,
-                                     self.host_alarm_id, ELabel.ON,
+                                     cls.host_alarm_id, ELabel.ON,
                                      host_to_add)
-                self.host_alarm_id += 1
+                cls.host_alarm_id += 1
 
             # Add Host Tests
             for j in range(num_of_tests_per_host):
-                add_connected_vertex(g, TEST, TEST_ON_HOST, self.host_test_id,
+                add_connected_vertex(g, TEST, TEST_ON_HOST, cls.host_test_id,
                                      ELabel.ON, host_to_add)
-                self.host_test_id += 1
+                cls.host_test_id += 1
 
             # Add Host Vms
             for j in range(num_of_vms_per_host):
                 vm_to_add = add_connected_vertex(g,
                                                  RESOURCE,
                                                  NOVA_INSTANCE_DATASOURCE,
-                                                 self.vm_id,
+                                                 cls.vm_id,
                                                  ELabel.CONTAINS,
                                                  host_to_add,
                                                  True)
-                self.vm_id += 1
+                cls.vm_id += 1
-                self.vms.append(vm_to_add)
+                cls.vms.append(vm_to_add)
 
                 # Add Instance Alarms
                 for k in range(num_of_alarms_per_vm):
                     add_connected_vertex(g, ALARM, ALARM_ON_VM,
-                                         self.vm_alarm_id, ELabel.ON,
+                                         cls.vm_alarm_id, ELabel.ON,
                                          vm_to_add)
-                    self.vm_alarm_id += 1
+                    cls.vm_alarm_id += 1
 
         end = time.time()
         LOG.debug('Graph creation took ' + str(end - start) +
@@ -205,5 +194,10 @@ class GraphTestBase(base.BaseTest):
             num_of_vms_per_host + num_of_hosts_per_node * \
             num_of_vms_per_host * num_of_alarms_per_vm + \
             num_of_tests_per_host * num_of_hosts_per_node
-        assert expected_graph_size == len(g), 'Graph size'
+        if not expected_graph_size == len(g):
+            raise VitrageError('Init failed, graph size unexpected {0} != {1}'
+                               .format(expected_graph_size, len(g)))
+        # cls.assertEqual(
+        #     expected_graph_size,
+        #     len(g), 'num of vertex node')
         return g

vitrage/tests/unit/graph/test_graph.py

@@ -31,7 +31,7 @@ LOG = logging.getLogger(__name__)
 class TestGraph(GraphTestBase):
 
     def test_graph(self):
-        g = create_graph('test_graph')
+        g = NXGraph('test_graph')
         self.assertEqual('test_graph', g.name, 'graph name')
         self.assertEqual(0, len(g), 'graph __len__')
 
@@ -59,7 +59,7 @@ class TestGraph(GraphTestBase):
                          'graph copy vertex unchanged after update')
 
     def test_vertex_crud(self):
-        g = create_graph('test_vertex_crud')
+        g = NXGraph('test_vertex_crud')
         g.add_vertex(v_node)
         v = g.get_vertex(v_node.vertex_id)
         self.assertEqual(v_node[VProps.ID], v[VProps.ID],
@@ -125,7 +125,7 @@ class TestGraph(GraphTestBase):
     def test_update_vertices(self):
 
         # Test Setup
-        g = create_graph('test_update_vertices')
+        g = NXGraph('test_update_vertices')
         g.add_vertex(v_node)
         v_node_copy = g.get_vertex(v_node.vertex_id)
         v_node_copy[VProps.NAME] = 'test_node'
@@ -147,7 +147,7 @@ class TestGraph(GraphTestBase):
         self.assertEqual('test_host', updated_v_host[VProps.NAME])
 
     def test_edge_crud(self):
-        g = create_graph('test_edge_crud')
+        g = NXGraph('test_edge_crud')
         g.add_vertex(v_node)
         g.add_vertex(v_host)
         g.add_edge(e_node_to_host)
@@ -256,7 +256,7 @@ class TestGraph(GraphTestBase):
             vitrage_id='kuku',
             entity_category=NOVA_HOST_DATASOURCE)
 
-        g = create_graph('test_neighbors')
+        g = NXGraph('test_neighbors')
         g.add_vertex(v1)
         g.add_vertex(v2)
         g.add_vertex(v3)
@@ -386,7 +386,7 @@ class TestGraph(GraphTestBase):
                                'Check neighbors for not connected vertex')
 
     def test_get_vertices(self):
-        g = create_graph('test_get_vertices')
+        g = NXGraph('test_get_vertices')
         g.add_vertex(v_node)
         g.add_vertex(v_host)
         g.add_edge(e_node_to_host)
@@ -436,7 +436,7 @@ class TestGraph(GraphTestBase):
     # noinspection PyAttributeOutsideInit
     def test_graph_callbacks(self):
 
-        g = create_graph('test_graph_callbacks')
+        g = NXGraph('test_graph_callbacks')
         self.result = None
 
         def callback(pre_item,
@@ -498,14 +498,14 @@ class TestGraph(GraphTestBase):
                                     target_id=v4.vertex_id,
                                     relationship_type='KUKU_v3_v4')
 
-        g1 = create_graph('test_union')
+        g1 = NXGraph('test_union')
         g1.add_vertex(v1)
         g1.add_vertex(v2)
         g1.add_vertex(v3)
         g1.add_edge(e_v1_v2)
         g1.add_edge(e_v2_v3)
 
-        g2 = create_graph('test_union_')
+        g2 = NXGraph('test_union_')
         g2.add_vertex(v3)
         g2.add_vertex(v4)
         g2.add_edge(e_v3_v4)

vitrage/tests/unit/graph/test_graph_algo.py

@@ -23,15 +23,32 @@ Tests for `vitrage` graph driver algorithms
 from vitrage.common.constants import EdgeLabel
 from vitrage.common.constants import EdgeProperties as EProps
 from vitrage.common.constants import VertexProperties as VProps
+from vitrage.graph.algo_driver.algorithm import Mapping
 from vitrage.graph.driver.elements import Edge
-from vitrage.graph import create_algorithm, Mapping, Direction  # noqa
+from vitrage.graph.driver.graph import Direction
 from vitrage.tests.unit.graph.base import *  # noqa
 
 
 class GraphAlgorithmTest(GraphTestBase):
 
+    # noinspection PyPep8Naming
+    @classmethod
+    def setUpClass(cls):
+        cls.vm_id = 10000000
+        cls.vm_alarm_id = 30000000
+        cls.vms = []
+        cls.host_alarm_id = 20000000
+        cls.host_test_id = 40000000
+        cls.entity_graph = cls._create_entity_graph(
+            'entity_graph',
+            num_of_hosts_per_node=ENTITY_GRAPH_HOSTS_PER_CLUSTER,
+            num_of_vms_per_host=ENTITY_GRAPH_VMS_PER_HOST,
+            num_of_alarms_per_host=ENTITY_GRAPH_ALARMS_PER_HOST,
+            num_of_alarms_per_vm=ENTITY_GRAPH_ALARMS_PER_VM,
+            num_of_tests_per_host=ENTITY_GRAPH_TESTS_PER_HOST)
+
     def test_graph_query_vertices(self):
-        ga = create_algorithm(self.entity_graph)
+        ga = self.entity_graph.algo
 
         query = {'==': {VProps.TYPE: OPENSTACK_CLUSTER}}
         subgraph = ga.graph_query_vertices(query)
@@ -186,11 +203,14 @@ class GraphAlgorithmTest(GraphTestBase):
                          'We filtered the ON relationship, so no alarms '
                          'should exist')
 
-    def test_no_match_graph_query_vertices(self):
+        # Undo changes made by this test
-        ga = create_algorithm(self.entity_graph)
+        host_instance_edge[VProps.IS_DELETED] = False
+        self.entity_graph.update_edge(host_instance_edge)
+        self.entity_graph.remove_edge(new_edge)
 
+    def test_no_match_graph_query_vertices(self):
         query = {'==': {VProps.TYPE: 'test'}}
-        subgraph = ga.graph_query_vertices(query)
+        subgraph = self.entity_graph.algo.graph_query_vertices(query)
         self.assertEqual(
             0,
             subgraph.num_vertices(), 'num of vertex node')
@@ -201,7 +221,7 @@ class GraphAlgorithmTest(GraphTestBase):
         Using the entity graph (created above) as a big graph we search
         for a sub graph match
         """
-        ga = create_algorithm(self.entity_graph)
+        ga = self.entity_graph.algo
 
         # Get ids of some of the elements in the entity graph:
         vm_alarm = self.entity_graph.get_vertex(
@@ -210,7 +230,7 @@ class GraphAlgorithmTest(GraphTestBase):
             ALARM_ON_HOST + str(self.host_alarm_id - 1))
 
         # Create a template for template matching
-        t = create_graph('template_graph')
+        t = NXGraph('template_graph')
         t_v_host_alarm = graph_utils.create_vertex(
             vitrage_id='1', entity_category=ALARM, entity_type=ALARM_ON_HOST)
         t_v_alarm_fail = graph_utils.create_vertex(

vitrage_tempest_tests/tests/api/base.py

@@ -29,8 +29,8 @@ from vitrage.datasources import NOVA_INSTANCE_DATASOURCE
 from vitrage.datasources import NOVA_ZONE_DATASOURCE
 from vitrage.datasources import OPENSTACK_CLUSTER
 from vitrage.datasources.static_physical import SWITCH
+from vitrage.graph.driver.networkx_graph import NXGraph
 from vitrage.graph import Edge
-from vitrage.graph import NXGraph
 from vitrage.graph import Vertex
 from vitrage import keystone_client
 from vitrage import os_clients