Examples with TestingRowConsumer io.crate.testing.TestingRowConsumer used on opensource projects

Example 61 with TestingRowConsumer

use of io.crate.testing.TestingRowConsumer in project crate by crate.

the class DistResultRXTaskTest method testNonSequentialBucketIds.

@Test
public void testNonSequentialBucketIds() throws Exception {
    TestingRowConsumer batchConsumer = new TestingRowConsumer();
    DistResultRXTask ctx = getPageDownstreamContext(batchConsumer, PassThroughPagingIterator.oneShot(), 3);
    PageBucketReceiver bucketReceiver = ctx.getBucketReceiver((byte) 0);
    assertThat(bucketReceiver, notNullValue());
    final PageResultListener mockListener = mock(PageResultListener.class);
    Bucket b1 = new CollectionBucket(Collections.singletonList(new Object[] { "foo" }));
    bucketReceiver.setBucket(0, b1, true, mockListener);
    Bucket b2 = new CollectionBucket(Collections.singletonList(new Object[] { "bar" }));
    bucketReceiver.setBucket(3, b2, true, mockListener);
    Bucket b3 = new CollectionBucket(Collections.singletonList(new Object[] { "universe" }));
    CheckPageResultListener checkPageResultListener = new CheckPageResultListener();
    bucketReceiver.setBucket(42, b3, false, checkPageResultListener);
    assertThat(checkPageResultListener.needMoreResult, is(true));
    bucketReceiver.setBucket(42, b3, true, checkPageResultListener);
    assertThat(checkPageResultListener.needMoreResult, is(false));
    List<Object[]> result = batchConsumer.getResult();
    assertThat(result.toArray(), arrayContainingInAnyOrder(new Object[] { "foo" }, new Object[] { "bar" }, new Object[] { "universe" }, new Object[] { "universe" }));
}

Example 62 with TestingRowConsumer

use of io.crate.testing.TestingRowConsumer in project crate by crate.

the class DistResultRXTaskTest method testListenerCalledAfterOthersHasFailed.

@Test
public void testListenerCalledAfterOthersHasFailed() throws Exception {
    TestingRowConsumer consumer = new TestingRowConsumer();
    DistResultRXTask ctx = getPageDownstreamContext(consumer, PassThroughPagingIterator.oneShot(), 2);
    PageBucketReceiver bucketReceiver = ctx.getBucketReceiver((byte) 0);
    assertThat(bucketReceiver, notNullValue());
    bucketReceiver.kill(new Exception("dummy"));
    PageResultListener listener = mock(PageResultListener.class);
    bucketReceiver.setBucket(1, Bucket.EMPTY, true, listener);
    verify(listener, times(1)).needMore(false);
}

Example 63 with TestingRowConsumer

use of io.crate.testing.TestingRowConsumer in project crate by crate.

the class DistResultRXTaskTest method testSetBucketOnAKilledCtxReleasesListener.

@Test
public void testSetBucketOnAKilledCtxReleasesListener() throws Exception {
    TestingRowConsumer consumer = new TestingRowConsumer();
    DistResultRXTask ctx = getPageDownstreamContext(consumer, PassThroughPagingIterator.oneShot(), 2);
    PageBucketReceiver bucketReceiver = ctx.getBucketReceiver((byte) 0);
    assertThat(bucketReceiver, notNullValue());
    ctx.kill(new InterruptedException("killed"));
    CompletableFuture<Void> listenerReleased = new CompletableFuture<>();
    bucketReceiver.setBucket(0, Bucket.EMPTY, false, needMore -> listenerReleased.complete(null));
    // Must not timeout
    listenerReleased.get(1, TimeUnit.SECONDS);
}

Example 64 with TestingRowConsumer

use of io.crate.testing.TestingRowConsumer in project crate by crate.

the class CountTaskTest method testClose.

@Test
public void testClose() throws Exception {
    CompletableFuture<Long> future = new CompletableFuture<>();
    CountOperation countOperation = mock(CountOperation.class);
    when(countOperation.count(eq(txnCtx), any(), any(Symbol.class))).thenReturn(future);
    CountTask countTask = new CountTask(countPhaseWithId(1), txnCtx, countOperation, new TestingRowConsumer(), null);
    countTask.start();
    future.complete(1L);
    assertTrue(countTask.isClosed());
    // assure that there was no exception
    countTask.completionFuture().get();
    // on error
    future = new CompletableFuture<>();
    when(countOperation.count(eq(txnCtx), any(), any(Symbol.class))).thenReturn(future);
    countTask = new CountTask(countPhaseWithId(2), txnCtx, countOperation, new TestingRowConsumer(), null);
    countTask.start();
    future.completeExceptionally(new UnhandledServerException("dummy"));
    assertTrue(countTask.isClosed());
    expectedException.expectCause(CauseMatcher.cause(UnhandledServerException.class));
    countTask.completionFuture().get();
}