@Test
public void testSuspendResume()
throws Exception
{
SuspendingExecutor suspendingExecutor = createSuspendingExecutor();
HiveSplitSource hiveSplitSource = new HiveSplitSource("test", 10, suspendingExecutor);
// almost fill the source
for (int i = 0; i < 9; i++) {
hiveSplitSource.addToQueue(new TestSplit(i));
assertEquals(hiveSplitSource.getOutstandingSplitCount(), i + 1);
assertFalse(suspendingExecutor.isSuspended());
}
// add one more split so the source is now full and verify that the executor is suspended
hiveSplitSource.addToQueue(new TestSplit(10));
assertEquals(hiveSplitSource.getOutstandingSplitCount(), 10);
assertTrue(suspendingExecutor.isSuspended());
// remove one split so the source is no longer full and verify the executor is resumed
assertEquals(hiveSplitSource.getNextBatch(1).size(), 1);
assertEquals(hiveSplitSource.getOutstandingSplitCount(), 9);
assertFalse(suspendingExecutor.isSuspended());
// add two more splits so the source is now full and verify that the executor is suspended
hiveSplitSource.addToQueue(new TestSplit(11));
hiveSplitSource.addToQueue(new TestSplit(12));
assertEquals(hiveSplitSource.getOutstandingSplitCount(), 11);
assertTrue(suspendingExecutor.isSuspended());
// remove two splits so the source is no longer full and verify the executor is resumed
assertEquals(hiveSplitSource.getNextBatch(2).size(), 2);
assertEquals(hiveSplitSource.getOutstandingSplitCount(), 9);
assertFalse(suspendingExecutor.isSuspended());
}