@Before public void before() { when(koalaIdFactory.buildPId(eq(nodeIdStr))).thenReturn(nodeId); doAnswer( new Answer<PId>() { @Override public PId answer(InvocationOnMock invocation) throws Throwable { String id = (String) invocation.getArguments()[0]; if (id.equals(Instance.getUrl("i-123"))) return id123; if (id.equals(Instance.getUrl("i-456"))) return id456; return null; } }) .when(koalaIdFactory) .buildId(isA(String.class)); when(piIdBuilder.getPIdForEc2AvailabilityZone(anyString())).thenReturn(id123).thenReturn(id456); when(apiApplicationManager.newMessageContext()) .thenReturn(instanceManagerMessageContext) .thenReturn(secondMessageContext); when(apiApplicationManager.getKoalaIdFactory()).thenReturn(koalaIdFactory); AvailabilityZones zones = new AvailabilityZones(); zones.addAvailabilityZone( new AvailabilityZone( AVAILABILITY_ZONE_1_NAME, AVAILABILITY_ZONE_1_CODE, REGION_ZONE_1_CODE, AVAILABILITY_ZONE_1_NAME)); when(apiApplicationManager.getAvailabilityZonesRecord()).thenReturn(zones); when(koalaIdFactory.getRegion()).thenReturn(REGION_ZONE_1_CODE); setupDht(); rebootInstanceServiceHelper.setPiIdBuilder(piIdBuilder); }