private void validateAllocatedContainers(boolean isReplicated, String containerTemplateLink)
        throws Throwable {
      QueryTask.Query kindClause =
          new QueryTask.Query()
              .setTermPropertyName(ServiceDocument.FIELD_NAME_KIND)
              .setTermMatchValue(Utils.buildKind(ContainerService.State.class));

      QueryTask.Query containerTemplateClause =
          new QueryTask.Query()
              .setTermPropertyName(
                  ContainerService.State.FIELD_NAME_CONTAINER_TEMPLATE_SERVICE_LINK)
              .setTermMatchValue(containerTemplateLink);

      QueryTask.QuerySpecification querySpecification = new QueryTask.QuerySpecification();
      querySpecification.query.addBooleanClause(kindClause);
      querySpecification.query.addBooleanClause(containerTemplateClause);

      QueryTask query = QueryTask.create(querySpecification).setDirect(true);

      NodeGroupBroadcastResponse queryResponse = testEnvironment.sendBroadcastQueryAndWait(query);
      Set<String> documentLinks = QueryTaskUtils.getBroadcastQueryResults(queryResponse);

      // Verify that count(replicas) == count(dockerVms) i.e. 1 container per vm
      int expectedReplicaCount = isReplicated ? dockerVms.size() : 1;
      assertThat(documentLinks.size(), is(expectedReplicaCount));

      // Verify that each container was assigned to a unique docker vm
      Set<String> uniqueVmLinks = new HashSet<>();
      for (String documentLink : documentLinks) {
        ContainerService.State state =
            testEnvironment.getServiceState(documentLink, ContainerService.State.class);
        uniqueVmLinks.add(state.vmServiceLink);
      }
      assertThat(uniqueVmLinks.size(), is(expectedReplicaCount));
    }
  /** Handle service periodic maintenance calls. */
  @Override
  public void handleMaintenance(Operation post) {
    post.complete();

    Operation.CompletionHandler handler =
        (Operation op, Throwable failure) -> {
          if (null != failure) {
            // query failed so abort and retry next time
            logFailure(failure);
            return;
          }

          NodeSelectorService.SelectOwnerResponse rsp =
              op.getBody(NodeSelectorService.SelectOwnerResponse.class);
          if (!getHost().getId().equals(rsp.ownerNodeId)) {
            ServiceUtils.logInfo(
                TaskTriggerService.this,
                "Host[%s]: Not owner of scheduler [%s] (Owner Info [%s])",
                getHost().getId(),
                getSelfLink(),
                Utils.toJson(rsp));
            return;
          }

          State state = new State();
          sendSelfPatch(state);
        };

    Operation selectOwnerOp =
        Operation.createPost(null)
            .setExpiration(ServiceUtils.computeExpirationTime(OWNER_SELECTION_TIMEOUT_MILLIS))
            .setCompletion(handler);
    getHost().selectOwner(null, getSelfLink(), selectOwnerOp);
  }
  /**
   * Interpret replication state.
   *
   * @param state
   * @return
   */
  private ReplicateImageStatus getReplicationStatus(ImageReplicatorService.State state) {
    ReplicateImageStatus result;
    switch (state.taskInfo.stage) {
      case CANCELLED:
        logger.error("Image replication failed: {}", Utils.toJson(state));
        result = new ReplicateImageStatus(ReplicateImageStatusCode.CANCELLED);
        result.setError("Image replication was cancelled.");
        break;

      case FAILED:
        if (ServiceStateUtils.isMinimumCopiesComplete(state, this.minReqCopies)) {
          result = new ReplicateImageStatus(ReplicateImageStatusCode.FINISHED);
        } else {
          logger.error("Image replication failed: {}", Utils.toJson(state));
          result = new ReplicateImageStatus(ReplicateImageStatusCode.FAILED);
          if (state.taskInfo != null && state.taskInfo.failure != null) {
            result.setError(
                String.format(
                    "Image replication failed. Error details: %s", state.taskInfo.failure.message));
          } else {
            result.setError("Image replication failed.");
          }
        }
        break;

      case FINISHED:
        result = new ReplicateImageStatus(ReplicateImageStatusCode.FINISHED);
        break;

      case STARTED:
        if (ServiceStateUtils.isMinimumCopiesComplete(state, this.minReqCopies)) {
          result = new ReplicateImageStatus(ReplicateImageStatusCode.FINISHED);
        } else {
          result = new ReplicateImageStatus(ReplicateImageStatusCode.IN_PROGRESS);
        }
        break;

      default:
        throw new RuntimeException(String.format("Unexpected stage %s.", state.taskInfo.stage));
    }
    return result;
  }
  /**
   * Get service state by serviceLink.
   *
   * @param path
   * @return
   * @throws Throwable
   */
  private ImageReplicatorService.State getLatestState(String path) throws Throwable {
    Operation getOperation =
        Operation.createGet(UriUtils.buildUri(dcpHost, path))
            .setReferer(UriUtils.buildUri(dcpHost, REFERRER_PATH))
            .setExpiration(Utils.getNowMicrosUtc() + dcpOperationTimeoutMicros)
            .setContextId(LoggingUtils.getRequestId());
    OperationLatch opLatch = new OperationLatch(getOperation);

    dcpHost.sendRequest(getOperation);
    return opLatch.await().getBody(ImageReplicatorService.State.class);
  }
  @VisibleForTesting
  protected State buildPatch(TaskState.TaskStage stage, @Nullable Throwable t) {
    State patchState = new State();
    patchState.taskState = new TaskState();
    patchState.taskState.stage = stage;

    if (null != t) {
      patchState.taskState.failure = Utils.toServiceErrorResponse(t);
    }

    return patchState;
  }
  /**
   * Does any additional processing after the patch operation has been completed.
   *
   * @param current
   */
  private void processPatch(final State current) {
    try {
      Type stateType = Class.forName(current.triggerStateClassName);
      ServiceDocument postState = Utils.fromJson(current.serializedTriggerState, stateType);
      postState.documentExpirationTimeMicros =
          ServiceUtils.computeExpirationTime(current.taskExpirationAgeMillis);

      Operation post =
          Operation.createPost(UriUtils.buildUri(getHost(), current.factoryServiceLink))
              .setBody(postState);
      this.sendRequest(post);
    } catch (ClassNotFoundException ex) {
      logFailure(ex);
    }
  }
  /**
   * Trigger image replication.
   *
   * @param request
   * @return
   * @throws Throwable
   */
  private String triggerReplication(ReplicateImageRequest request) throws Throwable {
    // Prepare replication service call.
    ImageReplicatorService.State postReq = new ImageReplicatorService.State();
    postReq.image = request.getImage();
    postReq.datastore = request.getDatastore();

    // Create the operation and call for replication.
    Operation postOperation =
        Operation.createPost(UriUtils.buildUri(dcpHost, ImageReplicatorServiceFactory.class))
            .setBody(postReq)
            .setReferer(UriUtils.buildUri(dcpHost, REFERRER_PATH))
            .setExpiration(Utils.getNowMicrosUtc() + dcpOperationTimeoutMicros)
            .setContextId(LoggingUtils.getRequestId());
    OperationLatch opLatch = new OperationLatch(postOperation);
    dcpHost.sendRequest(postOperation);
    Operation op = opLatch.await();

    // Return operation id.
    return op.getBody(ImageReplicatorService.State.class).documentSelfLink;
  }
    @Test
    public void testWaitForQuery() throws Throwable {
      ExampleService exampleService = new ExampleService();
      final ExampleService.ExampleServiceState exampleServiceState =
          new ExampleService.ExampleServiceState();
      exampleServiceState.name = UUID.randomUUID().toString();

      QueryTask.Query kindClause =
          new QueryTask.Query()
              .setTermPropertyName(ServiceDocument.FIELD_NAME_KIND)
              .setTermMatchValue(Utils.buildKind(ExampleService.ExampleServiceState.class));

      QueryTask.Query nameClause =
          new QueryTask.Query()
              .setTermPropertyName("name")
              .setTermMatchValue(exampleServiceState.name);

      QueryTask.QuerySpecification spec = new QueryTask.QuerySpecification();
      spec.query.addBooleanClause(kindClause);
      spec.query.addBooleanClause(nameClause);
      spec.options = EnumSet.of(QueryTask.QuerySpecification.QueryOption.EXPAND_CONTENT);

      QueryTask query = QueryTask.create(spec).setDirect(true);

      try {
        host.waitForQuery(
            query,
            new Predicate<QueryTask>() {
              @Override
              public boolean test(QueryTask queryTask) {
                return queryTask.results.documentLinks.size() > 0;
              }
            });

        Assert.fail("waitForQuery should not have succeeded before documents are created");
      } catch (RuntimeException runtimeException) {
        assertThat(
            runtimeException.getMessage(),
            is(equalToIgnoringCase("timeout waiting for query result.")));
      }

      Operation post = host.startServiceSynchronously(exampleService, exampleServiceState);
      assertThat(post.getStatusCode(), is(200));

      ExampleService.ExampleServiceState result =
          host.getServiceState(ExampleService.ExampleServiceState.class);
      assertThat(result.name, is(exampleServiceState.name));

      QueryTask response =
          host.waitForQuery(
              query,
              new Predicate<QueryTask>() {
                @Override
                public boolean test(QueryTask queryTask) {
                  return queryTask.results.documentLinks.size() > 0;
                }
              });
      assertThat(response.results.documentLinks.size(), is(1));

      // verify fields are passed down correctly
      for (Map.Entry<String, Object> document : response.results.documents.entrySet()) {
        ExampleService.ExampleServiceState docState =
            Utils.fromJson(document.getValue(), ExampleService.ExampleServiceState.class);
        assertThat(docState.name, is(equalTo(exampleServiceState.name)));
      }
    }