@Test  (timeout=5000)
public void testGetTask() throws IOException {

  AppContext appCtx = mock(AppContext.class);
  JobTokenSecretManager secret = mock(JobTokenSecretManager.class); 
  RMHeartbeatHandler rmHeartbeatHandler =
      mock(RMHeartbeatHandler.class);
  TaskHeartbeatHandler hbHandler = mock(TaskHeartbeatHandler.class);
  Dispatcher dispatcher = mock(Dispatcher.class);
  EventHandler ea = mock(EventHandler.class);
  when(dispatcher.getEventHandler()).thenReturn(ea);
  when(appCtx.getEventHandler()).thenReturn(ea);
  CheckpointAMPreemptionPolicy policy = new CheckpointAMPreemptionPolicy();
  policy.init(appCtx);
  MockTaskAttemptListenerImpl listener = 
    new MockTaskAttemptListenerImpl(appCtx, secret,
        rmHeartbeatHandler, hbHandler, policy);
  Configuration conf = new Configuration();
  listener.init(conf);
  listener.start();
  JVMId id = new JVMId("foo",1, true, 1);
  WrappedJvmID wid = new WrappedJvmID(id.getJobId(), id.isMap, id.getId());
  // Verify ask before registration.
  //The JVM ID has not been registered yet so we should kill it.
  JvmContext context = new JvmContext();
  context.jvmId = id; 
  JvmTask result = listener.getTask(context);
  assertNotNull(result);
  assertTrue(result.shouldDie);
  // Verify ask after registration but before launch. 
  // Don't kill, should be null.
  TaskAttemptId attemptID = mock(TaskAttemptId.class);
  Task task = mock(Task.class);
  //Now put a task with the ID
  listener.registerPendingTask(task, wid);
  result = listener.getTask(context);
  assertNull(result);
  // Unregister for more testing.
  listener.unregister(attemptID, wid);
  // Verify ask after registration and launch
  //Now put a task with the ID
  listener.registerPendingTask(task, wid);
  listener.registerLaunchedTask(attemptID, wid);
  verify(hbHandler).register(attemptID);
  result = listener.getTask(context);
  assertNotNull(result);
  assertFalse(result.shouldDie);
  // Don't unregister yet for more testing.
  //Verify that if we call it again a second time we are told to die.
  result = listener.getTask(context);
  assertNotNull(result);
  assertTrue(result.shouldDie);
  listener.unregister(attemptID, wid);
  // Verify after unregistration.
  result = listener.getTask(context);
  assertNotNull(result);
  assertTrue(result.shouldDie);
  listener.stop();
  // test JVMID
  JVMId jvmid = JVMId.forName("jvm_001_002_m_004");
  assertNotNull(jvmid);
  try {
    JVMId.forName("jvm_001_002_m_004_006");
    fail();
  } catch (IllegalArgumentException e) {
    assertEquals(e.getMessage(),
        "TaskId string : jvm_001_002_m_004_006 is not properly formed");
  }
}