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

Example 21 with TestingRowConsumer

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

the class ProjectionToProjectorVisitorTest method testAggregationProjector.

@Test
public void testAggregationProjector() throws Exception {
    AggregationProjection projection = new AggregationProjection(Arrays.asList(new Aggregation(avgSignature, avgSignature.getReturnType().createType(), Collections.singletonList(new InputColumn(1))), new Aggregation(CountAggregation.SIGNATURE, CountAggregation.SIGNATURE.getReturnType().createType(), Collections.singletonList(new InputColumn(0)))), RowGranularity.SHARD, AggregateMode.ITER_FINAL);
    Projector projector = visitor.create(projection, txnCtx, RamAccounting.NO_ACCOUNTING, memoryManager, UUID.randomUUID());
    assertThat(projector, instanceOf(AggregationPipe.class));
    BatchIterator<Row> batchIterator = projector.apply(InMemoryBatchIterator.of(new CollectionBucket(Arrays.asList($("foo", 10), $("bar", 20))), SENTINEL, true));
    TestingRowConsumer consumer = new TestingRowConsumer();
    consumer.accept(batchIterator, null);
    Bucket rows = consumer.getBucket();
    assertThat(rows.size(), is(1));
    assertThat(rows, contains(isRow(15.0, 2L)));
}

Example 22 with TestingRowConsumer

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

the class DistResultRXTaskTest method testPagingWithSortedPagingIterator.

@Test
public void testPagingWithSortedPagingIterator() throws Throwable {
    TestingRowConsumer batchConsumer = new TestingRowConsumer();
    DistResultRXTask ctx = getPageDownstreamContext(batchConsumer, new SortedPagingIterator<>(Comparator.comparingInt(r -> (int) r.get(0)), false), 2);
    Bucket b1 = new ArrayBucket(new Object[][] { new Object[] { 1 }, new Object[] { 1 } });
    Bucket b11 = new ArrayBucket(new Object[][] { new Object[] { 2 }, new Object[] { 2 } });
    PageBucketReceiver bucketReceiver = ctx.getBucketReceiver((byte) 0);
    assertThat(bucketReceiver, notNullValue());
    bucketReceiver.setBucket(0, b1, false, needMore -> {
        if (needMore) {
            bucketReceiver.setBucket(0, b11, true, mock(PageResultListener.class));
        }
    });
    Bucket b2 = new ArrayBucket(new Object[][] { new Object[] { 4 } });
    bucketReceiver.setBucket(1, b2, true, mock(PageResultListener.class));
    List<Object[]> result = batchConsumer.getResult();
    assertThat(TestingHelpers.printedTable(new CollectionBucket(result)), is("1\n" + "1\n" + "2\n" + "2\n" + "4\n"));
}

Example 23 with TestingRowConsumer

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

the class DistResultRXTaskTest method test_batch_iterator_is_completed_exceptionally_if_merge_buckets_on_next_page_fails.

@Test
public void test_batch_iterator_is_completed_exceptionally_if_merge_buckets_on_next_page_fails() throws Exception {
    TestingRowConsumer batchConsumer = new TestingRowConsumer();
    DistResultRXTask ctx = getPageDownstreamContext(batchConsumer, new FailOnMergePagingIterator<>(2), 2);
    PageBucketReceiver bucketReceiver = ctx.getBucketReceiver((byte) 0);
    assertThat(bucketReceiver, notNullValue());
    PageResultListener pageResultListener = mock(PageResultListener.class);
    Bucket bucket = new CollectionBucket(Collections.singletonList(new Object[] { "foo" }));
    bucketReceiver.setBucket(0, bucket, false, pageResultListener);
    bucketReceiver.setBucket(1, bucket, false, pageResultListener);
    bucketReceiver.setBucket(0, bucket, true, pageResultListener);
    bucketReceiver.setBucket(1, bucket, true, pageResultListener);
    expectedException.expect(RuntimeException.class);
    expectedException.expectMessage("raised on merge");
    batchConsumer.getResult();
}

Example 24 with TestingRowConsumer

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

the class DistResultRXTaskTest method test_batch_iterator_is_completed_exceptionally_if_first_merge_buckets_fails.

@Test
public void test_batch_iterator_is_completed_exceptionally_if_first_merge_buckets_fails() throws Exception {
    TestingRowConsumer batchConsumer = new TestingRowConsumer();
    DistResultRXTask ctx = getPageDownstreamContext(batchConsumer, new FailOnMergePagingIterator<>(1), 1);
    PageBucketReceiver bucketReceiver = ctx.getBucketReceiver((byte) 0);
    assertThat(bucketReceiver, notNullValue());
    PageResultListener pageResultListener = mock(PageResultListener.class);
    Bucket bucket = new CollectionBucket(Collections.singletonList(new Object[] { "foo" }));
    bucketReceiver.setBucket(0, bucket, true, pageResultListener);
    expectedException.expect(RuntimeException.class);
    expectedException.expectMessage("raised on merge");
    batchConsumer.getResult();
}

Example 25 with TestingRowConsumer

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

the class DistResultRXTaskTest method testListenersCalledWhenOtherUpstreamIsFailing.

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