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

Example 1 with TestingBatchConsumer

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

the class PageDownstreamContextTest method testListenerCalledAfterOthersHasFailed.

@Test
public void testListenerCalledAfterOthersHasFailed() throws Exception {
    TestingBatchConsumer consumer = new TestingBatchConsumer();
    PageDownstreamContext ctx = getPageDownstreamContext(consumer, PassThroughPagingIterator.oneShot(), 2);
    ctx.failure(0, new Exception("dummy"));
    PageResultListener listener = mock(PageResultListener.class);
    ctx.setBucket(1, Bucket.EMPTY, true, listener);
    verify(listener, times(1)).needMore(false);
}

Example 2 with TestingBatchConsumer

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

the class PageDownstreamContextTest method testKillCallsDownstream.

@Test
public void testKillCallsDownstream() throws Throwable {
    TestingBatchConsumer batchConsumer = new TestingBatchConsumer();
    PageDownstreamContext ctx = getPageDownstreamContext(batchConsumer, PassThroughPagingIterator.oneShot(), 3);
    final AtomicReference<Throwable> throwable = new AtomicReference<>();
    ctx.completionFuture().whenComplete((r, t) -> {
        if (t != null) {
            assertTrue(throwable.compareAndSet(null, t));
        } else {
            fail("Expected exception");
        }
    });
    ctx.kill(null);
    assertThat(throwable.get(), instanceOf(InterruptedException.class));
    expectedException.expect(InterruptedException.class);
    batchConsumer.getResult();
}

Example 3 with TestingBatchConsumer

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

the class CountContextTest method testClose.

@Test
public void testClose() throws Exception {
    SettableFuture<Long> future = SettableFuture.create();
    CountOperation countOperation = mock(CountOperation.class);
    when(countOperation.count(anyMap(), any(WhereClause.class))).thenReturn(future);
    CountContext countContext = new CountContext(1, countOperation, new TestingBatchConsumer(), null, WhereClause.MATCH_ALL);
    countContext.prepare();
    countContext.start();
    future.set(1L);
    assertTrue(countContext.future.closed());
    // assure that there was no exception
    countContext.completionFuture().get();
    // on error
    future = SettableFuture.create();
    when(countOperation.count(anyMap(), any(WhereClause.class))).thenReturn(future);
    countContext = new CountContext(2, countOperation, new TestingBatchConsumer(), null, WhereClause.MATCH_ALL);
    countContext.prepare();
    countContext.start();
    future.setException(new UnknownUpstreamFailure());
    assertTrue(countContext.future.closed());
    expectedException.expectCause(CauseMatcher.cause(UnknownUpstreamFailure.class));
    countContext.completionFuture().get();
}

Example 4 with TestingBatchConsumer

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

the class JobExecutionContextTest method testFailureClosesAllSubContexts.

@Test
public void testFailureClosesAllSubContexts() throws Exception {
    String localNodeId = "localNodeId";
    RoutedCollectPhase collectPhase = Mockito.mock(RoutedCollectPhase.class);
    Routing routing = Mockito.mock(Routing.class);
    when(routing.containsShards(localNodeId)).thenReturn(false);
    when(collectPhase.routing()).thenReturn(routing);
    when(collectPhase.maxRowGranularity()).thenReturn(RowGranularity.DOC);
    JobExecutionContext.Builder builder = new JobExecutionContext.Builder(UUID.randomUUID(), coordinatorNode, Collections.emptyList(), mock(JobsLogs.class));
    JobCollectContext jobCollectContext = new JobCollectContext(collectPhase, mock(MapSideDataCollectOperation.class), localNodeId, mock(RamAccountingContext.class), new TestingBatchConsumer(), mock(SharedShardContexts.class));
    TestingBatchConsumer batchConsumer = new TestingBatchConsumer();
    PageDownstreamContext pageDownstreamContext = spy(new PageDownstreamContext(Loggers.getLogger(PageDownstreamContext.class), "n1", 2, "dummy", batchConsumer, PassThroughPagingIterator.oneShot(), new Streamer[] { IntegerType.INSTANCE.streamer() }, mock(RamAccountingContext.class), 1));
    builder.addSubContext(jobCollectContext);
    builder.addSubContext(pageDownstreamContext);
    JobExecutionContext jobExecutionContext = builder.build();
    Exception failure = new Exception("failure!");
    jobCollectContext.close(failure);
    // other contexts must be killed with same failure
    verify(pageDownstreamContext, times(1)).innerKill(failure);
    final Field subContexts = JobExecutionContext.class.getDeclaredField("subContexts");
    subContexts.setAccessible(true);
    int size = ((ConcurrentMap<Integer, ExecutionSubContext>) subContexts.get(jobExecutionContext)).size();
    assertThat(size, is(0));
}

Example 5 with TestingBatchConsumer

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

the class PageDownstreamContextTest method testSetBucketOnAKilledCtxReleasesListener.

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