public synchronized void scheduleNopPollTask() { schedulePollTask(() -> { }); }
consumer.schedulePollTask(recordEnqueueTask); return consumer;
@Override public void run() { // add all the records with offset >= current partition position. int recordsAdded = 0; for (TopicPartition tp : assignedPartitions.get()) { long curPos = consumer.position(tp); for (ConsumerRecord<byte[], byte[]> r : records.get(tp)) { if (r.offset() >= curPos) { consumer.addRecord(r); recordsAdded++; } } } if (recordsAdded == 0) { if (config.get("inject.error.at.eof") != null) { consumer.setException(new KafkaException("Injected error in consumer.poll()")); } // MockConsumer.poll(timeout) does not actually wait even when there aren't any records. // Add a small wait here in order to avoid busy looping in the reader. Uninterruptibles.sleepUninterruptibly(10, TimeUnit.MILLISECONDS); //TODO: BEAM-4086: testUnboundedSourceWithoutBoundedWrapper() occasionally hangs // without this wait. Need to look into it. } consumer.schedulePollTask(this); } };
@Test public void testConsume(TestContext ctx) throws Exception { MockConsumer<String, String> mock = new MockConsumer<>(OffsetResetStrategy.EARLIEST); KafkaReadStream<String, String> consumer = createConsumer(vertx, mock); Async doneLatch = ctx.async(); consumer.handler(record -> { ctx.assertEquals("the_topic", record.topic()); ctx.assertEquals(0, record.partition()); ctx.assertEquals("abc", record.key()); ctx.assertEquals("def", record.value()); consumer.close(v -> doneLatch.complete()); }); consumer.subscribe(Collections.singleton("the_topic"), v -> { mock.schedulePollTask(() -> { mock.rebalance(Collections.singletonList(new TopicPartition("the_topic", 0))); mock.addRecord(new ConsumerRecord<>("the_topic", 0, 0L, "abc", "def")); mock.seek(new TopicPartition("the_topic", 0), 0L); }); }); }
@Test public void testBatch(TestContext ctx) throws Exception { int num = 50; MockConsumer<String, String> mock = new MockConsumer<>(OffsetResetStrategy.EARLIEST); KafkaReadStream<String, String> consumer = createConsumer(vertx, mock); Async doneLatch = ctx.async(); AtomicInteger count = new AtomicInteger(); consumer.handler(record -> { int val = count.getAndIncrement(); if (val < num) { ctx.assertEquals("the_topic", record.topic()); ctx.assertEquals(0, record.partition()); ctx.assertEquals("key-" + val, record.key()); ctx.assertEquals("value-" + val, record.value()); if (val == num - 1) { consumer.close(v -> doneLatch.complete()); } } }); consumer.subscribe(Collections.singleton("the_topic"), v -> { mock.schedulePollTask(() -> { mock.rebalance(Collections.singletonList(new TopicPartition("the_topic", 0))); mock.seek(new TopicPartition("the_topic", 0), 0); for (int i = 0; i < num; i++) { mock.addRecord(new ConsumerRecord<>("the_topic", 0, i, "key-" + i, "value-" + i)); } }); }); }
@Test public void testBatch(TestContext ctx) throws Exception { int num = 50; MockConsumer<String, String> mock = new MockConsumer<>(OffsetResetStrategy.EARLIEST); KafkaReadStream<String, String> consumer = createConsumer(vertx, mock); Async doneLatch = ctx.async(); AtomicInteger count = new AtomicInteger(); consumer.handler(record -> { int val = count.getAndIncrement(); if (val < num) { ctx.assertEquals("the_topic", record.topic()); ctx.assertEquals(0, record.partition()); ctx.assertEquals("key-" + val, record.key()); ctx.assertEquals("value-" + val, record.value()); if (val == num - 1) { consumer.close(v -> doneLatch.complete()); } } }); consumer.subscribe(Collections.singleton("the_topic"), v -> { mock.schedulePollTask(() -> { mock.rebalance(Collections.singletonList(new TopicPartition("the_topic", 0))); mock.seek(new TopicPartition("the_topic", 0), 0); for (int i = 0; i < num; i++) { mock.addRecord(new ConsumerRecord<>("the_topic", 0, i, "key-" + i, "value-" + i)); } }); }); }
@Test public void testConsume(TestContext ctx) throws Exception { MockConsumer<String, String> mock = new MockConsumer<>(OffsetResetStrategy.EARLIEST); KafkaReadStream<String, String> consumer = createConsumer(vertx, mock); Async doneLatch = ctx.async(); consumer.handler(record -> { ctx.assertEquals("the_topic", record.topic()); ctx.assertEquals(0, record.partition()); ctx.assertEquals("abc", record.key()); ctx.assertEquals("def", record.value()); consumer.close(v -> doneLatch.complete()); }); consumer.subscribe(Collections.singleton("the_topic"), v -> { mock.schedulePollTask(() -> { mock.rebalance(Collections.singletonList(new TopicPartition("the_topic", 0))); mock.addRecord(new ConsumerRecord<>("the_topic", 0, 0L, "abc", "def")); mock.seek(new TopicPartition("the_topic", 0), 0L); }); }); }
@Test public void testConsumeWithHeader(TestContext ctx) { MockConsumer<String, String> mock = new MockConsumer<>(OffsetResetStrategy.EARLIEST); KafkaReadStream<String, String> consumer = createConsumer(vertx, mock); Async doneLatch = ctx.async(); consumer.handler(record -> { ctx.assertEquals("the_topic", record.topic()); ctx.assertEquals(0, record.partition()); ctx.assertEquals("abc", record.key()); ctx.assertEquals("def", record.value()); Header[] headers = record.headers().toArray(); ctx.assertEquals(1, headers.length); Header header = headers[0]; ctx.assertEquals("header_key", header.key()); ctx.assertEquals("header_value", new String(header.value())); consumer.close(v -> doneLatch.complete()); }); consumer.subscribe(Collections.singleton("the_topic"), v -> { mock.schedulePollTask(() -> { mock.rebalance(Collections.singletonList(new TopicPartition("the_topic", 0))); mock.addRecord(new ConsumerRecord<>("the_topic", 0, 0L, 0L, TimestampType.NO_TIMESTAMP_TYPE, 0L, 0, 0, "abc", "def", new RecordHeaders(Collections.singletonList(new RecordHeader("header_key", "header_value".getBytes()))))); mock.seek(new TopicPartition("the_topic", 0), 0L); }); }); }
@Test public void testConsumeWithHeader(TestContext ctx) { MockConsumer<String, String> mock = new MockConsumer<>(OffsetResetStrategy.EARLIEST); KafkaReadStream<String, String> consumer = createConsumer(vertx, mock); Async doneLatch = ctx.async(); consumer.handler(record -> { ctx.assertEquals("the_topic", record.topic()); ctx.assertEquals(0, record.partition()); ctx.assertEquals("abc", record.key()); ctx.assertEquals("def", record.value()); Header[] headers = record.headers().toArray(); ctx.assertEquals(1, headers.length); Header header = headers[0]; ctx.assertEquals("header_key", header.key()); ctx.assertEquals("header_value", new String(header.value())); consumer.close(v -> doneLatch.complete()); }); consumer.subscribe(Collections.singleton("the_topic"), v -> { mock.schedulePollTask(() -> { mock.rebalance(Collections.singletonList(new TopicPartition("the_topic", 0))); mock.addRecord(new ConsumerRecord<>("the_topic", 0, 0L, 0L, TimestampType.NO_TIMESTAMP_TYPE, 0L, 0, 0, "abc", "def", new RecordHeaders(Collections.singletonList(new RecordHeader("header_key", "header_value".getBytes()))))); mock.seek(new TopicPartition("the_topic", 0), 0L); }); }); }