@SuppressWarnings("unchecked")
  private void setupInstancesInDht() {
    instance123 = setupInstanceWithReservationId("i-123", "r-123", ownerId);
    instance456 = setupInstanceWithReservationId("i-456", "r-123", ownerId);

    // setup dht reader to return the
    doAnswer(
            new Answer<Object>() {
              @Override
              public Object answer(InvocationOnMock invocation) throws Throwable {
                PiScatterGatherContinuation<Instance> piContinuation =
                    (PiScatterGatherContinuation<Instance>) invocation.getArguments()[1];

                if (invocation.getArguments()[0].equals(id123)) {
                  piContinuation.receiveResult(instance123);
                } else piContinuation.receiveResult(instance456);
                return null;
              }
            })
        .when(dhtReader)
        .getAnyAsync(or(eq(id123), eq(id456)), isA(PiScatterGatherContinuation.class));

    instanceIds = Arrays.asList(new String[] {"i-123", "i-456"});
    when(instance123.getNodeId()).thenReturn(nodeIdStr);
    when(instance456.getNodeId()).thenReturn(nodeIdStr);
  }
  @Test
  public void shouldRebootMultipleInstances() {
    // setup
    final Collection<String> instanceIds = Arrays.asList(new String[] {"i-123", "i-456"});
    when(instance123.getNodeId()).thenReturn(nodeIdStr);
    when(instance456.getNodeId()).thenReturn(nodeIdStr);
    when(piIdBuilder.getNodeIdFromNodeId(anyString())).thenReturn(nodeId);
    setupScatterGather(instanceIds);

    // act
    rebootInstanceServiceHelper.rebootInstances(ownerId, instanceIds, apiApplicationManager);

    // verify
    verify(instance123).setRestartRequested(eq(true));
    verify(instance456).setRestartRequested(eq(true));
    verify(instanceManagerMessageContext)
        .routePiMessageToApplication(
            eq(nodeId),
            eq(EntityMethod.UPDATE),
            eq(instance123),
            eq(InstanceManagerApplication.APPLICATION_NAME));
    verify(secondMessageContext)
        .routePiMessageToApplication(
            eq(nodeId),
            eq(EntityMethod.UPDATE),
            eq(instance456),
            eq(InstanceManagerApplication.APPLICATION_NAME));
  }
 @Override
 public Instance update(Instance existingEntity, Instance requestedEntity) {
   LOG.debug(
       String.format(
           "update(Existing Instance - %s,Requested Instance - %s)",
           existingEntity, requestedEntity));
   InstanceStateTransition transition = new InstanceStateTransition();
   Instance entityToReturn = null;
   if (existingEntity != null) {
     transition.setPreviousState(existingEntity.getState());
     // we set this here in case we can't write the state want.
     // this way the user doesn't get back a null. :D.
     transition.setNextState(existingEntity.getState());
     instanceStatusMap.put(existingEntity.getInstanceId(), transition);
     if (existingEntity.getNodeId() != null
         && existingEntity.getState() != InstanceState.SHUTTING_DOWN
         && existingEntity.getState() != InstanceState.TERMINATED) {
       existingEntity.setState(InstanceState.SHUTTING_DOWN);
       entityToReturn = existingEntity;
     } else if (existingEntity.getState() != InstanceState.TERMINATED) {
       LOG.info(
           String.format(
               "Instance %s is provisionally being terminated before an Instance Manager has claimed it by stamping a node id onto it",
               existingEntity.getInstanceId()));
       existingEntity.setState(InstanceState.TERMINATED);
       entityToReturn = existingEntity;
     }
   } else {
     LOG.warn(String.format("Expected non-null instance"));
   }
   return entityToReturn;
 }
  @Override
  public void handleResult(Instance result) {
    LOG.debug(String.format("handleResult(PiEntity - %s)", result));
    if (result != null) {
      if (result.getNodeId() != null) {
        LOG.debug(
            String.format(
                "Sending termination message for inst %s to node %s",
                result.getInstanceId(), result.getNodeId()));
        PId instanceRecordId = piIdBuilder.getNodeIdFromNodeId(result.getNodeId());
        MessageContext instanceMessageContext = messageContextFactory.newMessageContext();
        instanceMessageContext.routePiMessageToApplication(
            instanceRecordId,
            EntityMethod.DELETE,
            result,
            InstanceManagerApplication.APPLICATION_NAME);
      } else {
        LOG.debug(
            String.format(
                "Null node id detected, not sending termination message for %s to instance manager",
                result.getInstanceId()));
      }

      int instanceGlobalAvzCode =
          piIdBuilder.getGlobalAvailabilityZoneCodeFromEc2Id(result.getInstanceId());
      PId securityGroupId =
          piIdBuilder
              .getPId(SecurityGroup.getUrl(result.getUserId(), result.getSecurityGroupName()))
              .forGlobalAvailablityZoneCode(instanceGlobalAvzCode);
      LOG.debug(
          String.format(
              "Sending termination message for inst %s to network app node %s for sec group %s:%s",
              result.getInstanceId(),
              result.getUserId(),
              securityGroupId.toStringFull(),
              result.getSecurityGroupName()));
      MessageContext securityGroupMessageContext = messageContextFactory.newMessageContext();
      securityGroupMessageContext.routePiMessageToApplication(
          securityGroupId, EntityMethod.DELETE, result, NetworkManagerApplication.APPLICATION_NAME);
      InstanceStateTransition transition = instanceStatusMap.get(result.getInstanceId());
      if (transition != null) {
        transition.setNextState(result.getState());
      }
    } else {
      LOG.warn("Expected a non-null instance record");
    }
  }
  @Before
  public void setup() {
    when(piIdBuilder.getPIdForEc2AvailabilityZone(uri)).thenReturn(instancePId);
    when(piIdBuilder.getGlobalAvailabilityZoneCodeFromEc2Id(instanceId)).thenReturn(globalAvzCode);
    when(piIdBuilder.getPId(SecurityGroup.getUrl(userId, securityGroupName)))
        .thenReturn(securityGroupPId);
    when(piIdBuilder.getPiQueuePId(PiQueue.INSTANCE_NETWORK_MANAGER_TEARDOWN))
        .thenReturn(instanceNetworkManagerTeardownQueueId);
    when(securityGroupPId.forGlobalAvailablityZoneCode(globalAvzCode)).thenReturn(securityGroupPId);

    when(instanceNetworkManagerTeardownQueueId.forLocalScope(
            PiQueue.TERMINATE_INSTANCE.getNodeScope()))
        .thenReturn(instanceNetworkManagerTeardownQueueId);
    when(messageContextFactory.newMessageContext()).thenReturn(messageContext);
    when(dhtClientFactory.createReader()).thenReturn(dhtReader);

    when(resultInstance.getInstanceId()).thenReturn(instanceId);
    when(resultInstance.getNodeId()).thenReturn(instanceNodeId);
    when(resultInstance.getUserId()).thenReturn(userId);
    when(resultInstance.getSecurityGroupName()).thenReturn(securityGroupName);
  }