f4c21041bf7c8febdb1b14abc77efc96c1be139b
[demos/kafka/training] / src / test / java / de / juplo / kafka / ApplicationTests.java
1 package de.juplo.kafka;
2
3 import lombok.extern.slf4j.Slf4j;
4 import org.apache.kafka.clients.consumer.ConsumerRecord;
5 import org.apache.kafka.clients.consumer.KafkaConsumer;
6 import org.apache.kafka.clients.producer.KafkaProducer;
7 import org.apache.kafka.clients.producer.ProducerRecord;
8 import org.apache.kafka.common.TopicPartition;
9 import org.apache.kafka.common.serialization.*;
10 import org.apache.kafka.common.utils.Bytes;
11 import org.junit.jupiter.api.*;
12 import org.springframework.beans.factory.annotation.Autowired;
13 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
14 import org.springframework.boot.test.autoconfigure.data.mongo.AutoConfigureDataMongo;
15 import org.springframework.boot.test.context.ConfigDataApplicationContextInitializer;
16 import org.springframework.boot.test.context.TestConfiguration;
17 import org.springframework.context.annotation.Bean;
18 import org.springframework.context.annotation.Import;
19 import org.springframework.kafka.test.context.EmbeddedKafka;
20 import org.springframework.test.context.TestPropertySource;
21 import org.springframework.test.context.junit.jupiter.SpringJUnitConfig;
22
23 import java.time.Clock;
24 import java.time.Duration;
25 import java.util.*;
26 import java.util.concurrent.ExecutionException;
27 import java.util.concurrent.ExecutorService;
28 import java.util.function.BiConsumer;
29 import java.util.function.Function;
30 import java.util.stream.Collectors;
31 import java.util.stream.IntStream;
32
33 import static de.juplo.kafka.ApplicationTests.PARTITIONS;
34 import static de.juplo.kafka.ApplicationTests.TOPIC;
35 import static org.assertj.core.api.Assertions.*;
36 import static org.awaitility.Awaitility.*;
37
38
39 @SpringJUnitConfig(initializers = ConfigDataApplicationContextInitializer.class)
40 @TestMethodOrder(MethodOrderer.OrderAnnotation.class)
41 @TestPropertySource(
42                 properties = {
43                                 "consumer.bootstrap-server=${spring.embedded.kafka.brokers}",
44                                 "consumer.topic=" + TOPIC,
45                                 "consumer.commit-interval=1s",
46                                 "spring.mongodb.embedded.version=4.4.13" })
47 @EmbeddedKafka(topics = TOPIC, partitions = PARTITIONS)
48 @EnableAutoConfiguration
49 @AutoConfigureDataMongo
50 @Slf4j
51 class ApplicationTests
52 {
53         public static final String TOPIC = "FOO";
54         public static final int PARTITIONS = 10;
55
56
57         StringSerializer stringSerializer = new StringSerializer();
58
59         @Autowired
60         Serializer valueSerializer;
61         @Autowired
62         KafkaProducer<String, Bytes> kafkaProducer;
63         @Autowired
64         KafkaConsumer<String, String> kafkaConsumer;
65         @Autowired
66         PartitionStatisticsRepository partitionStatisticsRepository;
67         @Autowired
68         ApplicationProperties properties;
69         @Autowired
70         ExecutorService executor;
71         @Autowired
72         PartitionStatisticsRepository repository;
73         @Autowired
74         WordcountRebalanceListener wordcountRebalanceListener;
75         @Autowired
76         WordcountRecordHandler wordcountRecordHandler;
77
78         EndlessConsumer<String, String> endlessConsumer;
79         Map<TopicPartition, Long> oldOffsets;
80         Map<TopicPartition, Long> newOffsets;
81         Set<ConsumerRecord<String, String>> receivedRecords;
82
83
84         /** Tests methods */
85
86         @Test
87         void commitsCurrentOffsetsOnSuccess() throws ExecutionException, InterruptedException
88         {
89                 send100Messages(i ->  new Bytes(valueSerializer.serialize(TOPIC, i)));
90
91                 await("100 records received")
92                                 .atMost(Duration.ofSeconds(30))
93                                 .until(() -> receivedRecords.size() >= 100);
94
95                 await("Offsets committed")
96                                 .atMost(Duration.ofSeconds(10))
97                                 .untilAsserted(() ->
98                                 {
99                                         checkSeenOffsetsForProgress();
100                                         compareToCommitedOffsets(newOffsets);
101                                 });
102
103                 assertThatExceptionOfType(IllegalStateException.class)
104                                 .isThrownBy(() -> endlessConsumer.exitStatus())
105                                 .describedAs("Consumer should still be running");
106         }
107
108
109         /** Helper methods for the verification of expectations */
110
111         void compareToCommitedOffsets(Map<TopicPartition, Long> offsetsToCheck)
112         {
113                 doForCurrentOffsets((tp, offset) ->
114                 {
115                         Long expected = offsetsToCheck.get(tp) + 1;
116                         log.debug("Checking, if the offset for {} is {}", tp, expected);
117                         assertThat(offset)
118                                         .describedAs("Committed offset corresponds to the offset of the consumer")
119                                         .isEqualTo(expected);
120                 });
121         }
122
123         void checkSeenOffsetsForProgress()
124         {
125                 // Be sure, that some messages were consumed...!
126                 Set<TopicPartition> withProgress = new HashSet<>();
127                 partitions().forEach(tp ->
128                 {
129                         Long oldOffset = oldOffsets.get(tp);
130                         Long newOffset = newOffsets.get(tp);
131                         if (!oldOffset.equals(newOffset))
132                         {
133                                 log.debug("Progress for {}: {} -> {}", tp, oldOffset, newOffset);
134                                 withProgress.add(tp);
135                         }
136                 });
137                 assertThat(withProgress)
138                                 .describedAs("Some offsets must have changed, compared to the old offset-positions")
139                                 .isNotEmpty();
140         }
141
142
143         /** Helper methods for setting up and running the tests */
144
145         void doForCurrentOffsets(BiConsumer<TopicPartition, Long> consumer)
146         {
147                 partitions().forEach(tp ->
148                 {
149                         String partition = Integer.toString(tp.partition());
150                         Optional<Long> offset = partitionStatisticsRepository.findById(partition).map(document -> document.offset);
151                         consumer.accept(tp, offset.orElse(0l));
152                 });
153         }
154
155         List<TopicPartition> partitions()
156         {
157                 return
158                                 IntStream
159                                                 .range(0, PARTITIONS)
160                                                 .mapToObj(partition -> new TopicPartition(TOPIC, partition))
161                                                 .collect(Collectors.toList());
162         }
163
164
165         void send100Messages(Function<Long, Bytes> messageGenerator)
166         {
167                 long i = 0;
168
169                 for (int partition = 0; partition < 10; partition++)
170                 {
171                         for (int key = 0; key < 10; key++)
172                         {
173                                 Bytes value = messageGenerator.apply(++i);
174
175                                 ProducerRecord<String, Bytes> record =
176                                                 new ProducerRecord<>(
177                                                                 TOPIC,
178                                                                 partition,
179                                                                 Integer.toString(key%2),
180                                                                 value);
181
182                                 kafkaProducer.send(record, (metadata, e) ->
183                                 {
184                                         if (metadata != null)
185                                         {
186                                                 log.debug(
187                                                                 "{}|{} - {}={}",
188                                                                 metadata.partition(),
189                                                                 metadata.offset(),
190                                                                 record.key(),
191                                                                 record.value());
192                                         }
193                                         else
194                                         {
195                                                 log.warn(
196                                                                 "Exception for {}={}: {}",
197                                                                 record.key(),
198                                                                 record.value(),
199                                                                 e.toString());
200                                         }
201                                 });
202                         }
203                 }
204         }
205
206
207         @BeforeEach
208         public void init()
209         {
210                 oldOffsets = new HashMap<>();
211                 newOffsets = new HashMap<>();
212                 receivedRecords = new HashSet<>();
213
214                 doForCurrentOffsets((tp, offset) ->
215                 {
216                         oldOffsets.put(tp, offset - 1);
217                         newOffsets.put(tp, offset - 1);
218                 });
219
220                 TestRecordHandler<String, String> captureOffsetAndExecuteTestHandler =
221                                 new TestRecordHandler<String, String>(wordcountRecordHandler) {
222                                         @Override
223                                         public void onNewRecord(ConsumerRecord<String, String> record)
224                                         {
225                                                 newOffsets.put(
226                                                                 new TopicPartition(record.topic(), record.partition()),
227                                                                 record.offset());
228                                                 receivedRecords.add(record);
229                                         }
230                                 };
231
232                 endlessConsumer =
233                                 new EndlessConsumer<>(
234                                                 executor,
235                                                 properties.getClientId(),
236                                                 properties.getTopic(),
237                                                 kafkaConsumer,
238                                                 wordcountRebalanceListener,
239                                                 captureOffsetAndExecuteTestHandler);
240
241                 endlessConsumer.start();
242         }
243
244         @AfterEach
245         public void deinit()
246         {
247                 try
248                 {
249                         endlessConsumer.stop();
250                 }
251                 catch (Exception e)
252                 {
253                         log.info("Exception while stopping the consumer: {}", e.toString());
254                 }
255         }
256
257
258         @TestConfiguration
259         @Import(ApplicationConfiguration.class)
260         public static class Configuration
261         {
262                 @Bean
263                 Serializer<Long> serializer()
264                 {
265                         return new LongSerializer();
266                 }
267
268                 @Bean
269                 KafkaProducer<String, Bytes> kafkaProducer(ApplicationProperties properties)
270                 {
271                         Properties props = new Properties();
272                         props.put("bootstrap.servers", properties.getBootstrapServer());
273                         props.put("linger.ms", 100);
274                         props.put("key.serializer", StringSerializer.class.getName());
275                         props.put("value.serializer", BytesSerializer.class.getName());
276
277                         return new KafkaProducer<>(props);
278                 }
279         }
280 }