@Test
  public void testUpdatePartitionToFullAutoResource() throws Exception {
    assumeTrue(fullAutoCompatible);
    assumeTrue(dcStr.equals("DC1"));
    String clusterName = CLUSTER_NAME_PREFIX + CLUSTER_NAME_IN_STATIC_CLUSTER_MAP;
    Map<String, Map<String, Set<String>>> instancesUnderResourcesInDc = new HashMap<>();
    Map<String, Map<String, Set<Integer>>> partitionsUnderResourcesInDc = new HashMap<>();
    int basePort = 18088;
    int dataNode = 6;
    int basePartitionId = 0;
    int numPartition = 200;
    testHardwareLayout = constructInitialHardwareLayoutJSON(CLUSTER_NAME_IN_STATIC_CLUSTER_MAP, dataNode, basePort);
    testPartitionLayout = constructInitialPartitionLayoutJSON(testHardwareLayout, numPartition, null);
    testHardwareLayout.getHardwareLayout().getDatacenters().forEach(dc -> {
      instancesUnderResourcesInDc.put(dc.getName(), new HashMap<>());
      partitionsUnderResourcesInDc.put(dc.getName(), new HashMap<>());
    });
    long expectedResourceCount = 1;
    writeBootstrapOrUpgrade(expectedResourceCount, false, dataNode);
    addExpectedPartitionAndInstanceToResource(testPartitionLayout, FULL_AUTO_COMPATIBLE_RESOURCE_NAME_START_NUMBER,
        instancesUnderResourcesInDc, partitionsUnderResourcesInDc);
    verifyResourceInFullAuto(testHardwareLayout.getHardwareLayout(), instancesUnderResourcesInDc,
        partitionsUnderResourcesInDc);

    // Test 1: Add a new partition to resource when it hasn't turned on FULL AUTO
    basePartitionId += numPartition;
    basePort += dataNode * 2; // it has two datacenter
    List<DataNode> dataNodesForNewPartition = testHardwareLayout.getIndependentDataNodes(3);
    testPartitionLayout.addNewPartition(testHardwareLayout, dataNodesForNewPartition, DEFAULT_PARTITION_CLASS);
    writeBootstrapOrUpgrade(expectedResourceCount, true, dataNode);
    for (Datacenter dc : testHardwareLayout.getHardwareLayout().getDatacenters()) {
      String resourceName = String.valueOf(FULL_AUTO_COMPATIBLE_RESOURCE_NAME_START_NUMBER);
      partitionsUnderResourcesInDc.get(dc.getName()).get(resourceName).add(testPartitionLayout.getPartitionCount() - 1);
    }
    verifyResourceInFullAuto(testHardwareLayout.getHardwareLayout(), instancesUnderResourcesInDc,
        partitionsUnderResourcesInDc);

    // Test 2: Add a new group of partitions and instances and then turn second resource to FULL AUTO, then add a new partition
    basePartitionId += 1;
    int secondResourceBasePartitionId = basePartitionId;
    int secondResourceBasePort = basePort;
    TestHardwareLayout secondTestHardwareLayout =
        constructInitialHardwareLayoutJSON(CLUSTER_NAME_IN_STATIC_CLUSTER_MAP, dataNode, basePort);
    TestPartitionLayout secondTestPartitionLayout =
        constructInitialPartitionLayoutJSON(secondTestHardwareLayout, numPartition, null, basePartitionId);
    testPartitionLayout.merge(secondTestPartitionLayout);
    expectedResourceCount = 2;
    writeBootstrapOrUpgrade(expectedResourceCount, true, dataNode);
    addExpectedPartitionAndInstanceToResource(secondTestPartitionLayout,
        FULL_AUTO_COMPATIBLE_RESOURCE_NAME_START_NUMBER + 1, instancesUnderResourcesInDc, partitionsUnderResourcesInDc);

    turnOnFullAutoOnResource(FULL_AUTO_COMPATIBLE_RESOURCE_NAME_START_NUMBER + 1, dcStr, clusterName);
    // resource is in full auto now
    // Add another partition, but it will not be added to helix
    dataNodesForNewPartition = secondTestHardwareLayout.getIndependentDataNodes(3);
    testPartitionLayout.addNewPartition(testHardwareLayout, dataNodesForNewPartition, DEFAULT_PARTITION_CLASS);
    writeBootstrapOrUpgrade(expectedResourceCount, true, dataNode);
    verifyResourceInFullAuto(testHardwareLayout.getHardwareLayout(), instancesUnderResourcesInDc,
        partitionsUnderResourcesInDc);

    // Test 3: Remove a replica of a partition under a FULL AUTO resource
    Partition partition = (Partition) testPartitionLayout.getPartitionLayout()
        .getPartitions(null)
        .stream()
        .filter(partitionId -> partitionId.getId() >= secondResourceBasePartitionId)
        .findFirst()
        .get();
    List<ReplicaId> replicas = partition.getReplicas()
        .stream()
        .filter(replicaId -> replicaId.getDataNodeId().getDatacenterName().equals(dcStr))
        .collect(Collectors.toList());
    Set<String> instances = replicas.stream()
        .map(replicaId -> ClusterMapUtils.getInstanceName(replicaId.getDataNodeId()))
        .collect(Collectors.toSet());
    Replica replica = (Replica) replicas.get(0);
    partition.getReplicas().remove(replica);
    // No op
    writeBootstrapOrUpgrade(expectedResourceCount, true, dataNode);
    checkPartitionReplicas(dcStr, clusterName, FULL_AUTO_COMPATIBLE_RESOURCE_NAME_START_NUMBER + 1,
        partition.toPathString(), instances);

    // Test 4: Remove a replica of a partition not under FULL AUTO resource
    Partition partitionNotFullAuto = (Partition) testPartitionLayout.getPartitionLayout()
        .getPartitions(null)
        .stream()
        .filter(partitionId -> partitionId.getId() < secondResourceBasePartitionId)
        .findFirst()
        .get();
    List<ReplicaId> replicasNotFullAuto = partitionNotFullAuto.getReplicas()
        .stream()
        .filter(replicaId -> replicaId.getDataNodeId().getDatacenterName().equals(dcStr))
        .collect(Collectors.toList());
    Set<String> instancesNotFullAuto = replicasNotFullAuto.stream()
        .map(replicaId -> ClusterMapUtils.getInstanceName(replicaId.getDataNodeId()))
        .collect(Collectors.toSet());
    Replica replicaNotFullAuto = (Replica) replicasNotFullAuto.get(0);
    partitionNotFullAuto.getReplicas().remove(replicaNotFullAuto);
    try {
      writeBootstrapOrUpgrade(expectedResourceCount, true, dataNode);
      fail("Removing a replica from a FULL AUTO compatible resource should fail");
    } catch (Exception e) {
      partitionNotFullAuto.getReplicas().add(replicaNotFullAuto);
    }

    // Test 5: replace a replica of a partition under a FULL AUTO resource
    List<DataNode> dataNodeIdInFirstResource = testHardwareLayout.getAllDataNodesFromDc(dcStr)
        .stream()
        .filter(dn -> dn.getPort() < secondResourceBasePort)
        .collect(Collectors.toList());
    List<DataNode> dataNodeIdInSecondResource = testHardwareLayout.getAllDataNodesFromDc(dcStr)
        .stream()
        .filter(dn -> dn.getPort() >= secondResourceBasePort)
        .collect(Collectors.toList());

    DataNode dataNode1 = dataNodeIdInSecondResource.stream()
        .filter(dn -> !instances.contains(ClusterMapUtils.getInstanceName(dn)))
        .findFirst()
        .get();
    Disk disk = dataNode1.getDisks().get(0);
    testPartitionLayout.getPartitionLayout().addNewReplicas(partition, Collections.singletonList(disk));
    // No op
    writeBootstrapOrUpgrade(expectedResourceCount, true, dataNode);
    checkPartitionReplicas(dcStr, clusterName, FULL_AUTO_COMPATIBLE_RESOURCE_NAME_START_NUMBER + 1,
        partition.toPathString(), instances);

    // Test 6: replace a replica of a partition not under a FULL AUTO resource
    DataNode dataNode2 = dataNodeIdInFirstResource.stream()
        .filter(dn -> !instancesNotFullAuto.contains(ClusterMapUtils.getInstanceName(dn)))
        .findFirst()
        .get();
    Disk disk2 = dataNode2.getDisks().get(0);
    partitionNotFullAuto.getReplicas().remove(replicaNotFullAuto);
    testPartitionLayout.getPartitionLayout().addNewReplicas(partitionNotFullAuto, Collections.singletonList(disk2));
    writeBootstrapOrUpgrade(expectedResourceCount, true, dataNode);
    instancesNotFullAuto.add(ClusterMapUtils.getInstanceName(dataNode2));
    instancesNotFullAuto.remove(ClusterMapUtils.getInstanceName(replicaNotFullAuto.getDataNodeId()));
    checkPartitionReplicas(dcStr, clusterName, FULL_AUTO_COMPATIBLE_RESOURCE_NAME_START_NUMBER,
        partitionNotFullAuto.toPathString(), instancesNotFullAuto);

    // Test 7: remove a partition under a FULL AUTO resource
    testPartitionLayout.getPartitionLayout().removePartition(partition);
    writeBootstrapOrUpgrade(expectedResourceCount, true, dataNode);
    checkPartitionReplicas(dcStr, clusterName, FULL_AUTO_COMPATIBLE_RESOURCE_NAME_START_NUMBER + 1,
        partition.toPathString(), instances);

    // Test 8: replace a host under FULL AUTO resource
    DataNode randomDataNode = (DataNode) replicas.get(0).getDataNodeId();
    testPartitionLayout.replaceDataNodeWithNewOne(randomDataNode);
    writeBootstrapOrUpgrade(expectedResourceCount, true, dataNode);
    checkPartitionReplicas(dcStr, clusterName, FULL_AUTO_COMPATIBLE_RESOURCE_NAME_START_NUMBER + 1,
        partition.toPathString(), instances);

    // Test 9: replace a host not under a FULL AUTO resource
    randomDataNode = (DataNode) replicasNotFullAuto.get(0).getDataNodeId();
    final DataNode constDataNode = randomDataNode;
    Set<PartitionId> partitionBelongToDataNode = testPartitionLayout.getPartitionLayout()
        .getPartitions(null)
        .stream()
        .filter(partitionId -> partitionIncludeDataNode(partitionId, constDataNode))
        .collect(Collectors.toSet());
    DataNode newDataNode = testPartitionLayout.replaceDataNodeWithNewOne(randomDataNode);
    writeBootstrapOrUpgrade(expectedResourceCount, true, dataNode);
    checkInstanceInPartition(dcStr, clusterName, FULL_AUTO_COMPATIBLE_RESOURCE_NAME_START_NUMBER,
        partitionBelongToDataNode.stream().map(PartitionId::toPathString).collect(Collectors.toSet()),
        getInstanceName(randomDataNode), false);
    checkInstanceInPartition(dcStr, clusterName, FULL_AUTO_COMPATIBLE_RESOURCE_NAME_START_NUMBER,
        partitionBelongToDataNode.stream().map(PartitionId::toPathString).collect(Collectors.toSet()),
        getInstanceName(newDataNode), true);

    // Test 10: uplift resource, but only with half of the capacity and partitions
    basePort += dataNode * 2 + 2; // it has two datacenter
    basePartitionId += numPartition + 1; // it has two datacenter
    TestHardwareLayout thirdTestHardwareLayout =
        constructInitialHardwareLayoutJSON(CLUSTER_NAME_IN_STATIC_CLUSTER_MAP, dataNode / 2, basePort);
    TestPartitionLayout thirdTestPartitionLayout =
        constructInitialPartitionLayoutJSON(thirdTestHardwareLayout, numPartition / 2, null, basePartitionId);
    testPartitionLayout.merge(thirdTestPartitionLayout);
    expectedResourceCount = 3;
    writeBootstrapOrUpgrade(expectedResourceCount, true, dataNode);
    addExpectedPartitionAndInstanceToResource(secondTestPartitionLayout,
        FULL_AUTO_COMPATIBLE_RESOURCE_NAME_START_NUMBER + 2, instancesUnderResourcesInDc, partitionsUnderResourcesInDc);
    turnOnFullAutoOnResource(FULL_AUTO_COMPATIBLE_RESOURCE_NAME_START_NUMBER + 2, dcStr, clusterName);

    basePort += dataNode;
    basePartitionId += numPartition / 2;
    TestHardwareLayout fourthTestHardwareLayout =
        constructInitialHardwareLayoutJSON(CLUSTER_NAME_IN_STATIC_CLUSTER_MAP, dataNode / 2, basePort);
    TestPartitionLayout fourthTestPartitionLayout =
        constructInitialPartitionLayoutJSON(fourthTestHardwareLayout, numPartition / 2, null, basePartitionId);
    testPartitionLayout.merge(fourthTestPartitionLayout);
    // create a partition from third and fourth layout and add it to partition layout, it will block all the partitions from fourth layout
    List<DataNode> dataNodeIdInThirdResource = thirdTestHardwareLayout.getAllDataNodesFromDc(dcStr);
    List<DataNode> dataNodeIdInFourthResource = fourthTestHardwareLayout.getAllDataNodesFromDc(dcStr);
    // pick on data node from third and two from fourth
    DataNode firstDataNode = dataNodeIdInThirdResource.get(0);
    DataNode secondDataNode = dataNodeIdInFourthResource.get(0);
    DataNode lastDataNode = dataNodeIdInFourthResource.get(1);
    testPartitionLayout.addNewPartition(testHardwareLayout, Arrays.asList(firstDataNode, secondDataNode, lastDataNode),
        DEFAULT_PARTITION_CLASS);
    writeBootstrapOrUpgrade(expectedResourceCount, true, dataNode);
    addExpectedPartitionAndInstanceToResource(secondTestPartitionLayout,
        FULL_AUTO_COMPATIBLE_RESOURCE_NAME_START_NUMBER + 2, instancesUnderResourcesInDc, partitionsUnderResourcesInDc);

    // Test 11: replace a host not under a FULL AUTO resource, when the clusterconfig has topology aware turned on
    ConfigAccessor configAccessor = new ConfigAccessor("localhost:" + dcsToZkInfo.get(dcStr).getPort());
    ClusterConfig clusterConfig = configAccessor.getClusterConfig(clusterName);
    clusterConfig.setTopologyAwareEnabled(true);
    configAccessor.setClusterConfig(clusterName, clusterConfig);

    randomDataNode = (DataNode) replicasNotFullAuto.get(1).getDataNodeId();
    final DataNode constDataNode2 = randomDataNode;
    partitionBelongToDataNode = testPartitionLayout.getPartitionLayout()
        .getPartitions(null)
        .stream()
        .filter(partitionId -> partitionIncludeDataNode(partitionId, constDataNode2))
        .collect(Collectors.toSet());
    newDataNode = testPartitionLayout.replaceDataNodeWithNewOne(randomDataNode);
    writeBootstrapOrUpgrade(expectedResourceCount, true, dataNode);
    checkInstanceInPartition(dcStr, clusterName, FULL_AUTO_COMPATIBLE_RESOURCE_NAME_START_NUMBER,
        partitionBelongToDataNode.stream().map(PartitionId::toPathString).collect(Collectors.toSet()),
        getInstanceName(randomDataNode), false);
    checkInstanceInPartition(dcStr, clusterName, FULL_AUTO_COMPATIBLE_RESOURCE_NAME_START_NUMBER,
        partitionBelongToDataNode.stream().map(PartitionId::toPathString).collect(Collectors.toSet()),
        getInstanceName(newDataNode), true);

    InstanceConfig instanceConfig = configAccessor.getInstanceConfig(clusterName, getInstanceName(newDataNode));
    System.out.println("Domain for new data node: " + instanceConfig.getDomainAsString());
    Assert.assertNotNull(instanceConfig.getDomainAsString());
    Assert.assertFalse(instanceConfig.getDomainAsString().isEmpty());
    Assert.assertNotNull(instanceConfig.getInstanceCapacityMap());
    Assert.assertTrue(instanceConfig.getInstanceCapacityMap().containsKey("DISK"));
    configAccessor.close();
  }
Test update partitions that belong to FULL AUTO resource
@throws Exception
ambry-tools/src/integration-test/java/com/github/ambry/clustermap/HelixBootstrapUpgradeToolTest.java