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));
    }
 @AfterMethod
 public void tearDownTest() throws Throwable {
   if (sourceEnvironment != null) {
     sourceEnvironment.stop();
   }
   if (destinationEnvironment != null) {
     destinationEnvironment.stop();
   }
   apiClientFactory = null;
 }
    @Test(dataProvider = "SuccessCombinations")
    public void testEndToEndSuccess(boolean isReplicated, int existingContainers) throws Throwable {
      ContainerTemplateService.State templateStartState =
          TestHelper.getContainerTemplateServiceStartState();
      templateStartState.isReplicated = isReplicated;

      ContainerTemplateService.State containerTemplate =
          TestHelper.createContainerTemplateService(testEnvironment, templateStartState);

      VmService.State vmState = new VmService.State();

      for (int i = 0; i < existingContainers; i++) {
        vmState.documentSelfLink = dockerVms.get(i);
        TestHelper.createContainerService(testEnvironment, containerTemplate, vmState);
      }

      startState.containerTemplateDocumentLink = containerTemplate.documentSelfLink;
      startState.dockerVmDocumentLinks.clear();
      startState.dockerVmDocumentLinks.addAll(dockerVms);
      startState.singletonVmServiceLink = dockerVms.get(0);

      CreateContainerSpecTaskService.State finalState =
          testEnvironment.callServiceAndWaitForState(
              CreateContainerSpecTaskFactoryService.SELF_LINK,
              startState,
              CreateContainerSpecTaskService.State.class,
              (state) -> TaskUtils.finalTaskStages.contains(state.taskState.stage));

      TestHelper.assertTaskStateFinished(finalState.taskState);

      validateAllocatedContainers(isReplicated, containerTemplate.documentSelfLink);
    }
 @AfterMethod
 public void tearDownTest() throws Throwable {
   if (null != testEnvironment) {
     testEnvironment.stop();
     testEnvironment = null;
   }
 }
    @Test
    public void testSourceEnvironmentStoppedFailure() throws Throwable {
      createTestEnvironment();
      sourceEnvironment.stop();
      sourceEnvironment = null;

      InitializeDeploymentMigrationWorkflowService.State finalState =
          destinationEnvironment.callServiceAndWaitForState(
              InitializeDeploymentMigrationWorkflowFactoryService.SELF_LINK,
              startState,
              InitializeDeploymentMigrationWorkflowService.State.class,
              (state) -> TaskUtils.finalTaskStages.contains(state.taskState.stage));

      assertThat(
          finalState.taskState.stage,
          is(InitializeDeploymentMigrationWorkflowService.TaskState.TaskStage.FAILED));
    }
    private void setupContainer(String containerTemplateServiceLink, String vmServiceLink)
        throws Throwable {
      ContainerService.State container = new ContainerService.State();
      container.containerTemplateServiceLink = containerTemplateServiceLink;
      container.vmServiceLink = vmServiceLink;

      testEnvironment.callServiceSynchronously(
          ContainerFactoryService.SELF_LINK, container, ContainerService.State.class);
    }
    // TODO(hgadgil): This test does not run due to how DCP handles queries.
    // If a document link is specified and we try to retrieve it when the object does not exist, DCP
    // assumes that
    // the document is still being replicated so waits forever rather than returning failure
    // "NOT_FOUND".
    // Filed tracker story to enable this test when DCP enables the required behavior:
    // https://www.pivotaltracker.com/story/show/94686222
    @Test(enabled = false)
    public void testContainerTemplateNotFound() throws Throwable {
      CreateContainerSpecTaskService.State finalState =
          testEnvironment.callServiceAndWaitForState(
              CreateContainerSpecTaskFactoryService.SELF_LINK,
              startState,
              CreateContainerSpecTaskService.State.class,
              (state) -> TaskUtils.finalTaskStages.contains(state.taskState.stage));

      assertThat(finalState.taskState.stage, is(TaskState.TaskStage.FAILED));
    }
    private void createTestEnvironment() throws Throwable {
      ZookeeperClientFactory zkFactory = mock(ZookeeperClientFactory.class);
      sourceEnvironment =
          new TestEnvironment.Builder()
              .listeningExecutorService(listeningExecutorService)
              .apiClientFactory(apiClientFactory)
              .cloudServerSet(sourceCloudStore.getServerSet())
              .hostCount(1)
              .build();

      destinationEnvironment =
          new TestEnvironment.Builder()
              .hostCount(1)
              .apiClientFactory(apiClientFactory)
              .cloudServerSet(destinationCloudStore.getServerSet())
              .httpFileServiceClientFactory(httpFileServiceClientFactory)
              .zookeeperServersetBuilderFactory(zkFactory)
              .build();

      ZookeeperClient zkBuilder = mock(ZookeeperClient.class);
      doReturn(zkBuilder).when(zkFactory).create();
      doReturn(
              Collections.singleton(
                  new InetSocketAddress(
                      "127.0.0.1", sourceEnvironment.getHosts()[0].getState().httpPort)))
          .when(zkBuilder)
          .getServers(Matchers.startsWith("127.0.0.1:2181"), eq("cloudstore"));

      ServiceHost sourceHost = sourceEnvironment.getHosts()[0];
      startState.sourceLoadBalancerAddress = sourceHost.getPublicUri().toString();

      TestHelper.createHostService(sourceCloudStore, Collections.singleton(UsageTag.MGMT.name()));
      TestHelper.createHostService(sourceCloudStore, Collections.singleton(UsageTag.CLOUD.name()));
      DeploymentService.State deploymentService =
          TestHelper.createDeploymentService(destinationCloudStore);
      startState.destinationDeploymentId =
          ServiceUtils.getIDFromDocumentSelfLink(deploymentService.documentSelfLink);
    }
    @Test
    public void testSuccess() throws Throwable {
      createTestEnvironment();
      mockApiClient(true);
      MockHelper.mockHttpFileServiceClient(httpFileServiceClientFactory, true);

      InitializeDeploymentMigrationWorkflowService.State finalState =
          destinationEnvironment.callServiceAndWaitForState(
              InitializeDeploymentMigrationWorkflowFactoryService.SELF_LINK,
              startState,
              InitializeDeploymentMigrationWorkflowService.State.class,
              (state) -> TaskUtils.finalTaskStages.contains(state.taskState.stage));

      TestHelper.assertTaskStateFinished(finalState.taskState);
    }