1 files changed, 20 insertions, 7 deletions

diff --git a/aai-traversal/src/test/java/org/openecomp/aai/rest/search/QueryTest.java b/aai-traversal/src/test/java/org/openecomp/aai/rest/search/QueryTest.java
index 0df8ee3..2d40484 100644
--- a/aai-traversal/src/test/java/org/openecomp/aai/rest/search/QueryTest.java
+++ b/aai-traversal/src/test/java/org/openecomp/aai/rest/search/QueryTest.java
@@ -21,6 +21,8 @@
 package org.openecomp.aai.rest.search;
 
 import static org.junit.Assert.assertTrue;
+import static org.mockito.Matchers.any;
+import static org.mockito.Mockito.when;
 
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -31,20 +33,28 @@ import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversal;
 import org.apache.tinkerpop.gremlin.structure.Graph;
 import org.apache.tinkerpop.gremlin.structure.Vertex;
 import org.apache.tinkerpop.gremlin.tinkergraph.structure.TinkerGraph;
-
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 import org.openecomp.aai.exceptions.AAIException;
+import org.openecomp.aai.introspection.Loader;
+import org.openecomp.aai.introspection.LoaderFactory;
+import org.openecomp.aai.introspection.ModelType;
+import org.openecomp.aai.introspection.Version;
+import org.openecomp.aai.query.builder.GremlinTraversal;
 import org.openecomp.aai.serialization.db.EdgeRules;
 import org.openecomp.aai.serialization.db.exceptions.NoEdgeRuleFoundException;
+import org.openecomp.aai.serialization.engines.QueryStyle;
+import org.openecomp.aai.serialization.engines.TransactionalGraphEngine;
 
 public abstract class QueryTest {
 	
 	protected Graph graph;
 	private GremlinServerSingleton gremlinServerSingleton;
 	private GremlinGroovyShellSingleton shell;
-	
+	@Mock private TransactionalGraphEngine dbEngine;
 	protected final List<Vertex> expectedResult = new ArrayList<>();
 	protected final EdgeRules rules = EdgeRules.getInstance();
-	
+	protected Loader loader;
 	
 	public QueryTest() throws AAIException, NoEdgeRuleFoundException {
 		setUp();
@@ -52,22 +62,25 @@ public abstract class QueryTest {
 	public void setUp() throws AAIException, NoEdgeRuleFoundException {
 		System.setProperty("AJSC_HOME", ".");
 		System.setProperty("BUNDLECONFIG_DIR", "bundleconfig-local");
+		MockitoAnnotations.initMocks(this);
 		graph = TinkerGraph.open();
 		createGraph();
 		gremlinServerSingleton = GremlinServerSingleton.getInstance();
 		shell = GremlinGroovyShellSingleton.getInstance();
+		loader = LoaderFactory.createLoaderForVersion(ModelType.MOXY, Version.getLatest());
 	}
 	
 	public void run() {
 		
-		String query = "g." + gremlinServerSingleton.getStoredQuery(getQueryName());
-		
+		String query = gremlinServerSingleton.getStoredQuery(getQueryName());
 		Map<String, Object> params = new HashMap<>();
-		
+		addParam(params);
+		when(dbEngine.getQueryBuilder(any(QueryStyle.class))).thenReturn(new GremlinTraversal<>(loader, graph.traversal()));
+		query = GroovyQueryBuilderSingleton.getInstance().executeTraversal(dbEngine, query, params);
+		query = "g" + query;
 		GraphTraversal<Vertex, Vertex> g = graph.traversal().V();
 		addStartNode(g);
 		params.put("g", g);
-		addParam(params);
 		GraphTraversal<Vertex, Vertex> result = (GraphTraversal<Vertex, Vertex>)shell.executeTraversal(query, params);
 		
 		List<Vertex> vertices = result.toList();