1 files changed, 53 insertions, 1 deletions

diff --git a/models-interactions/model-actors/actor.test/src/main/java/org/onap/policy/controlloop/actor/test/BasicOperation.java b/models-interactions/model-actors/actor.test/src/main/java/org/onap/policy/controlloop/actor/test/BasicOperation.java
index f027c1c18..989f6a7c3 100644
--- a/models-interactions/model-actors/actor.test/src/main/java/org/onap/policy/controlloop/actor/test/BasicOperation.java
+++ b/models-interactions/model-actors/actor.test/src/main/java/org/onap/policy/controlloop/actor/test/BasicOperation.java
@@ -21,6 +21,8 @@
 package org.onap.policy.controlloop.actor.test;
 
 import static org.junit.Assert.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.when;
 
 import java.util.Map;
 import java.util.TreeMap;
@@ -29,6 +31,8 @@ import java.util.concurrent.CompletableFuture;
 import javax.ws.rs.core.Response;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.onap.policy.aai.AaiConstants;
+import org.onap.policy.aai.AaiCqResponse;
 import org.onap.policy.common.utils.coder.Coder;
 import org.onap.policy.common.utils.coder.CoderException;
 import org.onap.policy.common.utils.coder.StandardCoder;
@@ -36,9 +40,14 @@ import org.onap.policy.common.utils.resources.ResourceUtils;
 import org.onap.policy.common.utils.time.PseudoExecutor;
 import org.onap.policy.controlloop.VirtualControlLoopEvent;
 import org.onap.policy.controlloop.actorserviceprovider.ActorService;
+import org.onap.policy.controlloop.actorserviceprovider.Operation;
 import org.onap.policy.controlloop.actorserviceprovider.OperationOutcome;
+import org.onap.policy.controlloop.actorserviceprovider.Operator;
 import org.onap.policy.controlloop.actorserviceprovider.controlloop.ControlLoopEventContext;
+import org.onap.policy.controlloop.actorserviceprovider.impl.OperationPartial;
 import org.onap.policy.controlloop.actorserviceprovider.parameters.ControlLoopOperationParams;
+import org.onap.policy.controlloop.actorserviceprovider.spi.Actor;
+import org.onap.policy.controlloop.policy.PolicyResult;
 
 /**
  * Superclass for various Operation tests.
@@ -55,7 +64,22 @@ public class BasicOperation {
 
     @Mock
     protected ActorService service;
+    @Mock
+    protected Actor guardActor;
+    @Mock
+    protected Operator guardOperator;
+    @Mock
+    protected Operation guardOperation;
+    @Mock
+    protected Actor cqActor;
+    @Mock
+    protected Operator cqOperator;
+    @Mock
+    protected Operation cqOperation;
+    @Mock
+    protected AaiCqResponse cqResponse;
 
+    protected CompletableFuture<OperationOutcome> cqFuture;
     protected CompletableFuture<Response> future;
     protected ControlLoopOperationParams params;
     protected Map<String, String> enrichment;
@@ -89,12 +113,28 @@ public class BasicOperation {
     public void setUpBasic() {
         MockitoAnnotations.initMocks(this);
 
+        cqFuture = new CompletableFuture<>();
         future = new CompletableFuture<>();
 
         executor = new PseudoExecutor();
 
         makeContext();
 
+        when(service.getActor(OperationPartial.GUARD_ACTOR_NAME)).thenReturn(guardActor);
+        when(guardActor.getOperator(OperationPartial.GUARD_OPERATION_NAME)).thenReturn(guardOperator);
+        when(guardOperator.buildOperation(any())).thenReturn(guardOperation);
+
+        outcome = params.makeOutcome();
+        outcome.setResult(PolicyResult.SUCCESS);
+        when(guardOperation.start()).thenReturn(CompletableFuture.completedFuture(outcome));
+
+        when(service.getActor(AaiConstants.ACTOR_NAME)).thenReturn(cqActor);
+        when(cqActor.getOperator("CustomQuery")).thenReturn(cqOperator);
+        when(cqOperator.buildOperation(any())).thenReturn(cqOperation);
+
+        when(cqOperation.start()).thenReturn(cqFuture);
+
+        // get a fresh outcome
         outcome = params.makeOutcome();
     }
 
@@ -133,7 +173,7 @@ public class BasicOperation {
      *
      * @return payload data
      */
-    protected Map<String, String> makePayload() {
+    protected Map<String, Object> makePayload() {
         return null;
     }
 
@@ -162,4 +202,16 @@ public class BasicOperation {
 
         assertEquals(expected, json);
     }
+
+    /**
+     * Provides a response to a custom query.
+     *
+     * @param cq response to provide
+     */
+    protected void provideCqResponse(AaiCqResponse cq) {
+        context.setProperty(AaiCqResponse.CONTEXT_KEY, cq);
+        OperationOutcome outcome2 = params.makeOutcome();
+        outcome2.setResult(PolicyResult.SUCCESS);
+        cqFuture.complete(outcome2);
+    }
 }