@Test public void testFail() { MockSplitSource mockSource = new MockSplitSource() .setBatchSize(1) .increaseAvailableSplits(1) .atSplitCompletion(FAIL); try (SplitSource source = new BufferingSplitSource(mockSource, 100)) { assertFutureFailsWithMockFailure(getNextBatch(source, 2)); assertEquals(mockSource.getNextBatchInvocationCount(), 2); } }
public MockSplitSource atSplitCompletion(Action action) { atSplitDepletion = action; doGetNextBatch(); return this; }
@Override public ListenableFuture<SplitBatch> getNextBatch(ConnectorPartitionHandle partitionHandle, Lifespan lifespan, int maxSize) { if (partitionHandle != NOT_PARTITIONED) { throw new UnsupportedOperationException(); } checkArgument(Lifespan.taskWide().equals(lifespan)); checkState(nextBatchFuture.isDone(), "concurrent getNextBatch invocation"); nextBatchFuture = SettableFuture.create(); nextBatchMaxSize = maxSize; nextBatchInvocationCount++; doGetNextBatch(); return Futures.transform(nextBatchFuture, splits -> new SplitBatch(splits, isFinished()), directExecutor()); }
@Test public void testFailImmediate() { MockSplitSource mockSource = new MockSplitSource() .setBatchSize(1) .atSplitCompletion(FAIL); try (SplitSource source = new BufferingSplitSource(mockSource, 100)) { assertFutureFailsWithMockFailure(getNextBatch(source, 200)); assertEquals(mockSource.getNextBatchInvocationCount(), 1); } }
public MockSplitSource increaseAvailableSplits(int count) { checkArgument(atSplitDepletion == DO_NOTHING, "cannot increase available splits once split completion action is set"); totalSplits += count; doGetNextBatch(); return this; }
@Test public void testEmptySource() { MockSplitSource mockSource = new MockSplitSource() .setBatchSize(1) .atSplitCompletion(FINISH); try (SplitSource source = new BufferingSplitSource(mockSource, 100)) { requireFutureValue(getNextBatch(source, 200)) .assertSize(0) .assertNoMoreSplits(true); assertTrue(source.isFinished()); assertEquals(mockSource.getNextBatchInvocationCount(), 1); } }
@Test public void testFastSource() { MockSplitSource mockSource = new MockSplitSource() .setBatchSize(11) .increaseAvailableSplits(22) .atSplitCompletion(FINISH); try (SplitSource source = new BufferingSplitSource(mockSource, 10)) { requireFutureValue(getNextBatch(source, 200)) .assertSize(11) .assertNoMoreSplits(false); requireFutureValue(getNextBatch(source, 200)) .assertSize(11) .assertNoMoreSplits(true); assertTrue(source.isFinished()); assertEquals(mockSource.getNextBatchInvocationCount(), 2); } }
@Test public void testFinishedSetWithoutIndicationFromSplitBatch() { MockSplitSource mockSource = new MockSplitSource() .setBatchSize(1) .increaseAvailableSplits(1); try (SplitSource source = new BufferingSplitSource(mockSource, 100)) { requireFutureValue(getNextBatch(source, 1)) .assertSize(1) .assertNoMoreSplits(false); assertFalse(source.isFinished()); // Most of the time, mockSource.isFinished() returns the same value as // the SplitBatch.noMoreSplits field of the preceding mockSource.getNextBatch() call. // However, this is NOT always the case. // In this case, the preceding getNextBatch() indicates the noMoreSplits is false, // but the next isFinished call will return true. mockSource.atSplitCompletion(FINISH); requireFutureValue(getNextBatch(source, 1)) .assertSize(0) .assertNoMoreSplits(true); assertTrue(source.isFinished()); assertEquals(mockSource.getNextBatchInvocationCount(), 2); } }
@Test public void testSlowSource() { MockSplitSource mockSource = new MockSplitSource() .setBatchSize(1) .increaseAvailableSplits(25) .atSplitCompletion(FINISH); try (SplitSource source = new BufferingSplitSource(mockSource, 10)) { requireFutureValue(getNextBatch(source, 20)) .assertSize(10) .assertNoMoreSplits(false); requireFutureValue(getNextBatch(source, 6)) .assertSize(6) .assertNoMoreSplits(false); requireFutureValue(getNextBatch(source, 20)) .assertSize(9) .assertNoMoreSplits(true); assertTrue(source.isFinished()); assertEquals(mockSource.getNextBatchInvocationCount(), 25); } }
MockSplitSource mockSource = new MockSplitSource() .setBatchSize(1); try (SplitSource source = new BufferingSplitSource(mockSource, 10)) { mockSource.increaseAvailableSplits(9); assertFalse(nextBatchFuture.isDone()); mockSource.increaseAvailableSplits(1); requireFutureValue(nextBatchFuture) .assertSize(10) mockSource.atSplitCompletion(FINISH); requireFutureValue(nextBatchFuture) .assertSize(0) mockSource = new MockSplitSource() .setBatchSize(1); try (SplitSource source = new BufferingSplitSource(mockSource, 10)) { mockSource.increaseAvailableSplits(1); ListenableFuture<NextBatchResult> nextBatchFuture = getNextBatch(source, 10); assertFalse(nextBatchFuture.isDone()); mockSource.increaseAvailableSplits(9); requireFutureValue(nextBatchFuture) .assertSize(10) mockSource.increaseAvailableSplits(5); assertFalse(nextBatchFuture.isDone()); mockSource.atSplitCompletion(FINISH); requireFutureValue(nextBatchFuture)