| /* |
| * Licensed to the Apache Software Foundation (ASF) under one or more |
| * license agreements; and to You under the Apache License, version 2.0: |
| * |
| * https://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * This file is part of the Apache Pekko project, which was derived from Akka. |
| */ |
| |
| /* |
| * Copyright (C) 2014 - 2016 Softwaremill <https://softwaremill.com> |
| * Copyright (C) 2016 - 2020 Lightbend Inc. <https://www.lightbend.com> |
| */ |
| |
| package docs.scaladsl |
| |
| import org.apache.pekko |
| import pekko.Done |
| import pekko.kafka.Subscriptions |
| import pekko.kafka.scaladsl.{ Consumer, Producer, SpecBase } |
| import pekko.kafka.testkit.scaladsl.TestcontainersKafkaLike |
| import pekko.stream.scaladsl.{ Sink, Source } |
| import pekko.stream.testkit.scaladsl.StreamTestKit.assertAllStagesStopped |
| import org.apache.kafka.clients.producer.ProducerRecord |
| import org.apache.kafka.common.TopicPartition |
| |
| import scala.collection.immutable |
| |
| class AssignmentSpec extends SpecBase with TestcontainersKafkaLike { |
| |
| "subscription with partition assignment" must { |
| |
| "consume from the specified single topic" in assertAllStagesStopped { |
| val topic = createTopic() |
| val group = createGroupId() |
| val totalMessages = 100 |
| val producerCompletion = |
| Source(1 to totalMessages) |
| .map { msg => |
| new ProducerRecord(topic, 0, DefaultKey, msg.toString) |
| } |
| .runWith(Producer.plainSink(producerDefaults)) |
| |
| producerCompletion.futureValue |
| |
| // #single-topic |
| val subscription = Subscriptions.topics(topic) |
| val consumer = Consumer.plainSource(consumerDefaults.withGroupId(group), subscription) |
| // #single-topic |
| |
| val messages = |
| consumer.takeWhile(_.value().toInt < totalMessages, inclusive = true).runWith(Sink.seq) |
| messages.futureValue.size shouldBe totalMessages |
| } |
| |
| "consume from the specified topic pattern" in assertAllStagesStopped { |
| val suffix = (System.currentTimeMillis() % 10000).toInt |
| |
| val topics = immutable.Seq(createTopic(suffix), createTopic(suffix)) |
| val group = createGroupId() |
| val totalMessages = 100 |
| val producerCompletion = |
| Source(1 to totalMessages) |
| .mapConcat { msg => |
| topics.map(topic => new ProducerRecord(topic, 0, DefaultKey, msg.toString)) |
| } |
| .runWith(Producer.plainSink(producerDefaults)) |
| |
| producerCompletion.futureValue shouldBe Done |
| |
| // #topic-pattern |
| val pattern = s"topic-$suffix-[0-9]+" |
| val subscription = Subscriptions.topicPattern(pattern) |
| val consumer = Consumer.plainSource(consumerDefaults.withGroupId(group), subscription) |
| // #topic-pattern |
| |
| val messages = |
| consumer |
| .take(totalMessages * topics.length.toLong) |
| .runWith(Sink.seq) |
| val received = messages.futureValue |
| if (received.size != totalMessages * topics.size) { |
| received.foreach(r => println(s"topic=${r.topic()} value=${r.value()}")) |
| } |
| received.size shouldBe totalMessages * topics.size |
| } |
| |
| "consume from the specified partition" in assertAllStagesStopped { |
| val topic = createTopic(suffix = 0, partitions = 2) |
| val totalMessages = 100 |
| val producerCompletion = |
| Source(1 to totalMessages) |
| .map { msg => |
| val partition = msg % 2 |
| new ProducerRecord(topic, partition, DefaultKey, msg.toString) |
| } |
| .runWith(Producer.plainSink(producerDefaults)) |
| |
| producerCompletion.futureValue |
| |
| // #assingment-single-partition |
| val partition = 0 |
| val subscription = Subscriptions.assignment(new TopicPartition(topic, partition)) |
| val consumer = Consumer.plainSource(consumerDefaults, subscription) |
| // #assingment-single-partition |
| |
| val messages = consumer.take(totalMessages.toLong / 2).map(_.value().toInt).runWith(Sink.seq) |
| messages.futureValue.map(_ % 2 shouldBe 0) |
| } |
| |
| "consume from the specified partition and offset" in assertAllStagesStopped { |
| val topic = createTopic() |
| val totalMessages = 100 |
| val producerCompletion = |
| Source(1 to totalMessages) |
| .map { msg => |
| new ProducerRecord(topic, 0, DefaultKey, msg.toString) |
| } |
| .runWith(Producer.plainSink(producerDefaults)) |
| |
| producerCompletion.futureValue |
| |
| // #assingment-single-partition-offset |
| val partition = 0 |
| val offset: Long = totalMessages.toLong / 2 |
| val subscription = Subscriptions.assignmentWithOffset(new TopicPartition(topic, partition) -> offset) |
| val consumer = Consumer.plainSource(consumerDefaults, subscription) |
| // #assingment-single-partition-offset |
| |
| val messages = consumer.take(totalMessages.toLong / 2).map(_.offset()).runWith(Sink.seq) |
| messages.futureValue.map(_ - offset).zipWithIndex.map { case (offs, idx) => offs - idx }.sum shouldBe 0 |
| } |
| |
| "consume from the specified partition and timestamp" in assertAllStagesStopped { |
| val topic = createTopic() |
| val totalMessages = 100 |
| val producerCompletion = |
| Source(1 to totalMessages) |
| .map { msg => |
| new ProducerRecord(topic, 0, System.currentTimeMillis(), DefaultKey, msg.toString) |
| } |
| .runWith(Producer.plainSink(producerDefaults)) |
| |
| producerCompletion.futureValue |
| |
| // #assingment-single-partition-timestamp |
| val partition = 0 |
| val now = System.currentTimeMillis |
| val messagesSince: Long = now - 5000 |
| val subscription = Subscriptions.assignmentOffsetsForTimes(new TopicPartition(topic, partition) -> messagesSince) |
| val consumer = Consumer.plainSource(consumerDefaults, subscription) |
| // #assingment-single-partition-timestamp |
| |
| val messages = |
| consumer.takeWhile(_.value().toInt < totalMessages, inclusive = true).map(_.timestamp()).runWith(Sink.seq) |
| messages.futureValue.map(_ - now).count(_ > 5000) shouldBe 0 |
| } |
| } |
| |
| "illegal subscriptions" should { |
| "fail the stream with invalid topic" in { |
| val streamCompletion = Consumer |
| .plainSource(consumerDefaults.withGroupId(createGroupId()), Subscriptions.topics("illegal topic name")) |
| .runWith(Sink.ignore) |
| streamCompletion.failed.futureValue shouldBe a[org.apache.kafka.common.errors.InvalidTopicException] |
| } |
| |
| "fail the stream for empty subscriptions list" in { |
| val streamCompletion = Consumer |
| .plainSource(consumerDefaults.withGroupId(createGroupId()), Subscriptions.topics()) |
| .runWith(Sink.ignore) |
| streamCompletion.failed.futureValue shouldBe a[java.lang.IllegalStateException] |
| } |
| |
| "fail the stream for `null` topic name" in { |
| val streamCompletion = Consumer |
| .plainSource(consumerDefaults.withGroupId(createGroupId()), Subscriptions.topics(Set(null: String))) |
| .runWith(Sink.ignore) |
| streamCompletion.failed.futureValue shouldBe a[java.lang.IllegalArgumentException] |
| } |
| |
| "fail the stream with invalid topic for assignments" in { |
| val streamCompletion = Consumer |
| .plainSource(consumerDefaults.withGroupId(createGroupId()), |
| Subscriptions.assignment(new TopicPartition("illegal topic name", 2))) |
| .runWith(Sink.ignore) |
| streamCompletion.failed.futureValue shouldBe a[org.apache.kafka.common.errors.InvalidTopicException] |
| } |
| |
| "fail the stream for illegal topic pattern" in { |
| val streamCompletion = Consumer |
| .plainSource(consumerDefaults.withGroupId(createGroupId()), Subscriptions.topicPattern("illegal regex (*")) |
| .runWith(Sink.ignore) |
| streamCompletion.failed.futureValue shouldBe a[java.util.regex.PatternSyntaxException] |
| } |
| |
| "fail the stream for illegal offsets for times" in { |
| val streamCompletion = Consumer |
| .plainSource( |
| consumerDefaults.withGroupId(createGroupId()), |
| Subscriptions.assignmentOffsetsForTimes(new TopicPartition("topic", 0) -> 232L, |
| new TopicPartition("topic", 1) -> -232L)) |
| .runWith(Sink.ignore) |
| streamCompletion.failed.futureValue shouldBe a[java.lang.IllegalArgumentException] |
| } |
| } |
| |
| // #testkit |
| |
| } |
| // #testkit |
