Fix kafka streams latest dep test (#13546)

Author: laurit

instrumentation/kafka/kafka-streams-0.11/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/kafkastreams/KafkaStreamsBaseTest.java

@@ -28,6 +28,7 @@
 import org.apache.kafka.clients.admin.NewTopic;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.clients.producer.KafkaProducer;
 import org.apache.kafka.clients.producer.Producer;
@@ -146,7 +147,7 @@ static void awaitUntilConsumerIsReady() throws InterruptedException {
       return;
     }
     for (int i = 0; i < 10; i++) {
-      consumer.poll(0);
+      poll(Duration.ofMillis(100));
       if (consumerReady.await(1, TimeUnit.SECONDS)) {
         break;
       }
@@ -157,6 +158,10 @@ static void awaitUntilConsumerIsReady() throws InterruptedException {
     consumer.seekToBeginning(Collections.emptyList());
   }
 
+  public static ConsumerRecords<Integer, String> poll(Duration duration) {
+    return KafkaStreamsReflectionUtil.poll(consumer, duration);
+  }
+
   static Context getContext(Headers headers) {
     String traceparent =
         new String(

instrumentation/kafka/kafka-streams-0.11/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/kafkastreams/KafkaStreamsDefaultTest.java

@@ -53,7 +53,7 @@ class KafkaStreamsDefaultTest extends KafkaStreamsBaseTest {
   @Test
   void testKafkaProduceAndConsumeWithStreamsInBetween() throws Exception {
     Properties config = new Properties();
-    config.putAll(producerProps(KafkaStreamsBaseTest.kafka.getBootstrapServers()));
+    config.putAll(producerProps(kafka.getBootstrapServers()));
     config.put(StreamsConfig.APPLICATION_ID_CONFIG, "test-application");
     config.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.Integer().getClass().getName());
     config.put(
@@ -74,12 +74,11 @@ void testKafkaProduceAndConsumeWithStreamsInBetween() throws Exception {
     streams.start();
 
     String greeting = "TESTING TESTING 123!";
-    KafkaStreamsBaseTest.producer.send(new ProducerRecord<>(STREAM_PENDING, 10, greeting));
+    producer.send(new ProducerRecord<>(STREAM_PENDING, 10, greeting));
 
     awaitUntilConsumerIsReady();
     @SuppressWarnings("PreferJavaTimeOverload")
-    ConsumerRecords<Integer, String> records =
-        KafkaStreamsBaseTest.consumer.poll(Duration.ofSeconds(10).toMillis());
+    ConsumerRecords<Integer, String> records = poll(Duration.ofSeconds(10));
     Headers receivedHeaders = null;
     for (ConsumerRecord<Integer, String> record : records) {
       Span.current().setAttribute("testing", 123);

instrumentation/kafka/kafka-streams-0.11/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/kafkastreams/KafkaStreamsReflectionUtil.java

@@ -7,7 +7,10 @@
 
 import java.lang.reflect.Constructor;
 import java.lang.reflect.Method;
+import java.time.Duration;
 import java.util.Properties;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.common.serialization.Serde;
 import org.apache.kafka.common.serialization.Serdes;
 import org.apache.kafka.streams.KafkaStreams;
@@ -18,6 +21,32 @@
  * streams.
  */
 class KafkaStreamsReflectionUtil {
+  private static final Method consumerPollDurationMethod = getConsumerPollMethod(Duration.class);
+  private static final Method consumerPollLongMethod = getConsumerPollMethod(long.class);
+
+  private static Method getConsumerPollMethod(Class<?>... types) {
+    try {
+      return Consumer.class.getMethod("poll", types);
+    } catch (NoSuchMethodException e) {
+      return null;
+    }
+  }
+
+  @SuppressWarnings("unchecked")
+  public static <K, V> ConsumerRecords<K, V> poll(Consumer<K, V> consumer, Duration duration) {
+    try {
+      if (consumerPollDurationMethod != null) {
+        // not present in early versions
+        return (ConsumerRecords<K, V>) consumerPollDurationMethod.invoke(consumer, duration);
+      } else if (consumerPollLongMethod != null) {
+        // not present in 4.x
+        return (ConsumerRecords<K, V>) consumerPollLongMethod.invoke(consumer, duration.toMillis());
+      }
+    } catch (Exception exception) {
+      throw new IllegalStateException(exception);
+    }
+    throw new IllegalStateException("poll method not found");
+  }
 
   private KafkaStreamsReflectionUtil() {}
 

instrumentation/kafka/kafka-streams-0.11/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/kafkastreams/KafkaStreamsSuppressReceiveSpansTest.java

@@ -50,7 +50,7 @@ class KafkaStreamsSuppressReceiveSpansTest extends KafkaStreamsBaseTest {
   @Test
   void testKafkaProduceAndConsumeWithStreamsInBetween() throws Exception {
     Properties config = new Properties();
-    config.putAll(producerProps(KafkaStreamsBaseTest.kafka.getBootstrapServers()));
+    config.putAll(producerProps(kafka.getBootstrapServers()));
     config.put(StreamsConfig.APPLICATION_ID_CONFIG, "test-application");
     config.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.Integer().getClass().getName());
     config.put(
@@ -71,12 +71,10 @@ void testKafkaProduceAndConsumeWithStreamsInBetween() throws Exception {
     streams.start();
 
     String greeting = "TESTING TESTING 123!";
-    KafkaStreamsBaseTest.producer.send(new ProducerRecord<>(STREAM_PENDING, 10, greeting));
+    producer.send(new ProducerRecord<>(STREAM_PENDING, 10, greeting));
 
     // check that the message was received
-    @SuppressWarnings("PreferJavaTimeOverload")
-    ConsumerRecords<Integer, String> records =
-        KafkaStreamsBaseTest.consumer.poll(Duration.ofSeconds(10).toMillis());
+    ConsumerRecords<Integer, String> records = poll(Duration.ofSeconds(10));
     Headers receivedHeaders = null;
     for (ConsumerRecord<Integer, String> record : records) {
       Span.current().setAttribute("testing", 123);