@Test
public void testNotifyContainerCompleted() throws IOException {

  AgentProviderService aps = createAgentProviderService(new Configuration());
  AgentProviderService mockAps = Mockito.spy(aps);
  doNothing().when(mockAps).publishApplicationInstanceData(anyString(), anyString(), anyCollection());
  ContainerId cid = new MockContainerId(1);
  String id = cid.toString();
  ContainerId cid2 = new MockContainerId(2);
  mockAps.getAllocatedPorts().put("a", "100");
  mockAps.getAllocatedPorts(id).put("a", "100");
  mockAps.getAllocatedPorts(id).put("b", "101");
  mockAps.getAllocatedPorts("cid2").put("c", "102");
  mockAps.getComponentInstanceData().put("cid2", new HashMap<String, String>());
  mockAps.getComponentInstanceData().put(id, new HashMap<String, String>());
  mockAps.getComponentStatuses().put("cid2_HM", new ComponentInstanceState("HM", cid2, "aid"));
  mockAps.getComponentStatuses().put(id + "_HM", new ComponentInstanceState("HM", cid, "aid"));
  Assert.assertNotNull(mockAps.getComponentInstanceData().get(id));
  Assert.assertNotNull(mockAps.getComponentInstanceData().get("cid2"));
  Assert.assertNotNull(mockAps.getComponentStatuses().get(id + "_HM"));
  Assert.assertNotNull(mockAps.getComponentStatuses().get("cid2_HM"));
  Assert.assertEquals(mockAps.getAllocatedPorts().size(), 1);
  Assert.assertEquals(mockAps.getAllocatedPorts(id).size(), 2);
  Assert.assertEquals(mockAps.getAllocatedPorts("cid2").size(), 1);
  // Make the call
  mockAps.notifyContainerCompleted(new MockContainerId(1));
  Assert.assertEquals(mockAps.getAllocatedPorts().size(), 0);
  Assert.assertEquals(mockAps.getAllocatedPorts(id).size(), 0);
  Assert.assertEquals(mockAps.getAllocatedPorts("cid2").size(), 1);
  Assert.assertNull(mockAps.getComponentInstanceData().get(id));
  Assert.assertNotNull(mockAps.getComponentInstanceData().get("cid2"));
  Assert.assertNull(mockAps.getComponentStatuses().get(id + "_HM"));
  Assert.assertNotNull(mockAps.getComponentStatuses().get("cid2_HM"));
}