@Override
      public TaskInfo getNextTask() throws InterruptedException {

        while (true) {
          lock.lock();
          try {
            while (!shouldRun) {
              triggerRunCondition.await();
            }
            // Preven subsequent runs until a new trigger is set.
            shouldRun = false;
          } finally {
            lock.unlock();
          }
          TaskInfo taskInfo = delayedTaskQueue.peek();
          if (taskInfo == null) {
            LOG.info("Triggered getTask but the queue is empty");
            lastState = STATE_NULL_FOUND;
            signalRunComplete();
            continue;
          }
          if (taskInfo.shouldDelayForLocality(
              LlapTaskSchedulerServiceForTestControlled.this.clock.getTime())) {
            LOG.info("Triggered getTask but the first element is not ready to execute");
            lastState = STATE_TIMEOUT_NOT_EXPIRED;
            signalRunComplete();
            continue;
          } else {
            delayedTaskQueue.poll(); // Remove the previously peeked element.
            lastState = STATE_RETURNED_TASK;
            return taskInfo;
          }
        }
      }
Note: this only works for testing the lack of invocations from the main thread.
llap-tez/src/test/org/apache/hadoop/hive/llap/tezplugins/TestLlapTaskSchedulerService.java