@Test
  public void testGetAndSetReplicaSealedState() {
    //setup HelixParticipant and dependencies
    ClusterMapConfig clusterMapConfig = new ClusterMapConfig(new VerifiableProperties(props));
    String instanceName = ClusterMapUtils.getInstanceName("localhost", clusterMapConfig.clusterMapPort);
    HelixParticipant helixParticipant =
        new HelixParticipant(mock(HelixClusterManager.class), clusterMapConfig, new HelixFactory(),
            new MetricRegistry(), getDefaultZkConnectStr(clusterMapConfig), true);
    ZKHelixAdmin helixAdmin = new ZKHelixAdmin("localhost:" + zkInfo.getPort());
    DataNodeConfig dataNodeConfig = getDataNodeConfigInHelix(helixAdmin, instanceName);

    Set<String> localPartitionNames = new HashSet<>();
    dataNodeConfig.getDiskConfigs()
        .values()
        .forEach(diskConfig -> localPartitionNames.addAll(diskConfig.getReplicaConfigs().keySet()));
    String partitionIdStr = localPartitionNames.iterator().next();
    String partitionIdStr2 = localPartitionNames.stream().filter(p -> !p.equals(partitionIdStr)).findFirst().get();
    ReplicaId replicaId = createMockAmbryReplica(partitionIdStr);
    ReplicaId replicaId2 = createMockAmbryReplica(partitionIdStr2);

    //Make sure the current sealedReplicas list is empty
    List<String> sealedReplicas = helixParticipant.getSealedReplicas();
    List<String> partiallySealedReplicas = helixParticipant.getPartiallySealedReplicas();
    assertEquals("sealedReplicas should be empty", Collections.emptyList(), sealedReplicas);
    assertEquals("partiallySealedReplicas should be empty", Collections.emptyList(), partiallySealedReplicas);

    String sealedListName = "sealedReplicas";
    String partiallySealedListName = "partiallySealedReplicas";

    //Check that invoking setReplicaSealedState with a non-AmbryReplica ReplicaId throws an IllegalArgumentException
    ReplicaId notAmbryReplica = createMockNotAmbryReplica(partitionIdStr);
    try {
      helixParticipant.setReplicaSealedState(notAmbryReplica, ReplicaSealStatus.SEALED);
      fail("Expected an IllegalArgumentException here");
    } catch (IllegalArgumentException e) {
      //Expected exception
    }
    try {
      helixParticipant.setReplicaSealedState(notAmbryReplica, ReplicaSealStatus.PARTIALLY_SEALED);
      fail("Expected an IllegalArgumentException here");
    } catch (IllegalArgumentException e) {
      //Expected exception
    }

    //Check that invoking setReplicaSealedState adds the partition to the list of sealed replicas
    helixParticipant.setReplicaSealedState(replicaId, ReplicaSealStatus.SEALED);
    sealedReplicas = helixParticipant.getSealedReplicas();
    partiallySealedReplicas = helixParticipant.getPartiallySealedReplicas();
    listIsExpectedSize(sealedReplicas, 1, sealedListName);
    listIsExpectedSize(partiallySealedReplicas, 0, partiallySealedListName);
    assertTrue(sealedReplicas.contains(partitionIdStr));

    //Seal another replicaId
    helixParticipant.setReplicaSealedState(replicaId2, ReplicaSealStatus.SEALED);
    sealedReplicas = helixParticipant.getSealedReplicas();
    partiallySealedReplicas = helixParticipant.getPartiallySealedReplicas();
    listIsExpectedSize(sealedReplicas, 2, sealedListName);
    listIsExpectedSize(partiallySealedReplicas, 0, partiallySealedListName);
    assertTrue(sealedReplicas.contains(partitionIdStr2));
    assertTrue(sealedReplicas.contains(partitionIdStr));

    //Check that sealed replica list doesn't take duplicates (and that dups are detected by partitionId comparison, not
    //replicaId object comparison
    ReplicaId dup = createMockAmbryReplica(partitionIdStr);
    helixParticipant.setReplicaSealedState(dup, ReplicaSealStatus.SEALED);
    helixParticipant.setReplicaSealedState(replicaId2, ReplicaSealStatus.SEALED);
    sealedReplicas = helixParticipant.getSealedReplicas();
    partiallySealedReplicas = helixParticipant.getPartiallySealedReplicas();
    listIsExpectedSize(sealedReplicas, 2, sealedListName);
    listIsExpectedSize(partiallySealedReplicas, 0, partiallySealedListName);
    assertTrue(sealedReplicas.contains(partitionIdStr2));
    assertTrue(sealedReplicas.contains(partitionIdStr));

    //Check that invoking setReplicaSealedState with ReplicaSealStatus.NOT_SEALED removes partition from list of sealed replicas
    helixParticipant.setReplicaSealedState(replicaId, ReplicaSealStatus.NOT_SEALED);
    sealedReplicas = helixParticipant.getSealedReplicas();
    partiallySealedReplicas = helixParticipant.getPartiallySealedReplicas();
    listIsExpectedSize(sealedReplicas, 1, sealedListName);
    listIsExpectedSize(partiallySealedReplicas, 0, partiallySealedListName);
    assertTrue(sealedReplicas.contains(partitionIdStr2));
    assertFalse(sealedReplicas.contains(partitionIdStr));

    //Unsealing a replicaId that's already been unsealed doesn't hurt anything
    helixParticipant.setReplicaSealedState(replicaId, ReplicaSealStatus.NOT_SEALED);
    sealedReplicas = helixParticipant.getSealedReplicas();
    partiallySealedReplicas = helixParticipant.getPartiallySealedReplicas();
    listIsExpectedSize(sealedReplicas, 1, sealedListName);
    listIsExpectedSize(partiallySealedReplicas, 0, partiallySealedListName);

    //Unsealing all replicas yields expected behavior (and unseal works by partitionId, not replicaId itself)
    dup = createMockAmbryReplica(partitionIdStr2);
    helixParticipant.setReplicaSealedState(dup, ReplicaSealStatus.NOT_SEALED);
    sealedReplicas = helixParticipant.getSealedReplicas();
    partiallySealedReplicas = helixParticipant.getPartiallySealedReplicas();
    listIsExpectedSize(sealedReplicas, 0, sealedListName);
    listIsExpectedSize(partiallySealedReplicas, 0, partiallySealedListName);

    //Check that invoking setReplicaSealedState adds the partition to the list of partially sealed replicas
    helixParticipant.setReplicaSealedState(replicaId, ReplicaSealStatus.PARTIALLY_SEALED);
    sealedReplicas = helixParticipant.getSealedReplicas();
    partiallySealedReplicas = helixParticipant.getPartiallySealedReplicas();
    listIsExpectedSize(sealedReplicas, 0, sealedListName);
    listIsExpectedSize(partiallySealedReplicas, 1, partiallySealedListName);
    assertTrue(partiallySealedReplicas.contains(partitionIdStr));

    //Partially Seal another replicaId
    helixParticipant.setReplicaSealedState(replicaId2, ReplicaSealStatus.PARTIALLY_SEALED);
    sealedReplicas = helixParticipant.getSealedReplicas();
    partiallySealedReplicas = helixParticipant.getPartiallySealedReplicas();
    listIsExpectedSize(sealedReplicas, 0, sealedListName);
    listIsExpectedSize(partiallySealedReplicas, 2, partiallySealedListName);
    assertTrue(partiallySealedReplicas.contains(partitionIdStr2));
    assertTrue(partiallySealedReplicas.contains(partitionIdStr));

    //Seal one partially sealed replicaId
    helixParticipant.setReplicaSealedState(replicaId, ReplicaSealStatus.SEALED);
    sealedReplicas = helixParticipant.getSealedReplicas();
    partiallySealedReplicas = helixParticipant.getPartiallySealedReplicas();
    listIsExpectedSize(sealedReplicas, 1, sealedListName);
    listIsExpectedSize(partiallySealedReplicas, 1, partiallySealedListName);
    assertTrue(sealedReplicas.contains(partitionIdStr));
    assertTrue(partiallySealedReplicas.contains(partitionIdStr2));

    //Unseal one partially sealed replicaId
    helixParticipant.setReplicaSealedState(replicaId2, ReplicaSealStatus.NOT_SEALED);
    sealedReplicas = helixParticipant.getSealedReplicas();
    partiallySealedReplicas = helixParticipant.getPartiallySealedReplicas();
    listIsExpectedSize(sealedReplicas, 1, sealedListName);
    listIsExpectedSize(partiallySealedReplicas, 0, partiallySealedListName);
    assertTrue(sealedReplicas.contains(partitionIdStr));

    // Partially seal a sealed replicaId
    helixParticipant.setReplicaSealedState(replicaId, ReplicaSealStatus.PARTIALLY_SEALED);
    sealedReplicas = helixParticipant.getSealedReplicas();
    partiallySealedReplicas = helixParticipant.getPartiallySealedReplicas();
    listIsExpectedSize(sealedReplicas, 0, sealedListName);
    listIsExpectedSize(partiallySealedReplicas, 1, partiallySealedListName);
    assertTrue(partiallySealedReplicas.contains(partitionIdStr));

    helixAdmin.close();
    helixParticipant.close();
  }
Tests setReplicaSealedState method for {@link HelixParticipant}
@throws Exception
ambry-clustermap/src/test/java/com/github/ambry/clustermap/HelixParticipantTest.java