content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.javaagent.instrumentation.kafkaclients.v0_11;
  
  import static org.assertj.core.api.Assertions.assertThat;
  
  import io.opentelemetry.api.trace.SpanKind;
  import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.KafkaClientBaseTest;
  import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.KafkaClientPropagationBaseTest;
  import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
  import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
  import java.nio.charset.StandardCharsets;
  import java.time.Duration;
  import java.util.List;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.TimeoutException;
  import org.apache.kafka.clients.consumer.ConsumerRecord;
  import org.apache.kafka.clients.consumer.ConsumerRecords;
  import org.apache.kafka.clients.producer.ProducerRecord;
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.extension.RegisterExtension;
  
  class KafkaClientSuppressReceiveSpansTest extends KafkaClientPropagationBaseTest {
    @RegisterExtension
    static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
  
    @Test
    void testKafkaProduceAndConsume() throws InterruptedException {
      String greeting = "Hello Kafka!";
      testing.runWithSpan(
          "parent",
          () -> {
            ProducerRecord<Integer, String> producerRecord =
                new ProducerRecord<>(SHARED_TOPIC, 10, greeting);
            producerRecord
                .headers()
                // adding baggage header in w3c baggage format
                .add(
                    "baggage",
                    "test-baggage-key-1=test-baggage-value-1".getBytes(StandardCharsets.UTF_8))
                .add(
                    "baggage",
                    "test-baggage-key-2=test-baggage-value-2".getBytes(StandardCharsets.UTF_8));
            producer.send(
                producerRecord,
                (meta, ex) -> {
                  if (ex == null) {
                    testing.runWithSpan("producer callback", () -> {});
                  } else {
                    testing.runWithSpan("producer exception: " + ex, () -> {});
                  }
                });
          });
  
      awaitUntilConsumerIsReady();
      // check that the message was received
      ConsumerRecords<?, ?> records = poll(Duration.ofSeconds(5));
      for (ConsumerRecord<?, ?> record : records) {
        testing.runWithSpan(
            "processing",
            () -> {
              assertThat(record.key()).isEqualTo(10);
              assertThat(record.value()).isEqualTo(greeting);
            });
      }
  
      testing.waitAndAssertTraces(
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
                  span ->
                      span.hasName(SHARED_TOPIC + " publish")
                          .hasKind(SpanKind.PRODUCER)
                          .hasParent(trace.getSpan(0))
                          .hasAttributesSatisfyingExactly(sendAttributes("10", greeting, false)),
                  span ->
                      span.hasName(SHARED_TOPIC + " process")
                          .hasKind(SpanKind.CONSUMER)
                          .hasParent(trace.getSpan(1))
                          .hasAttributesSatisfyingExactly(
                              processAttributes("10", greeting, false, true)),
                  span ->
                      span.hasName("processing")
                          .hasKind(SpanKind.INTERNAL)
                          .hasParent(trace.getSpan(2)),
                  span ->
                      span.hasName("producer callback")
                          .hasKind(SpanKind.INTERNAL)
                          .hasParent(trace.getSpan(0))));
    }
  
    @Test
    void testPassThroughTombstone()
        throws ExecutionException, InterruptedException, TimeoutException {
      producer.send(new ProducerRecord<>(SHARED_TOPIC, null)).get(5, TimeUnit.SECONDS);
      awaitUntilConsumerIsReady();
      ConsumerRecords<?, ?> records = poll(Duration.ofSeconds(5));
      assertThat(records.count()).isEqualTo(1);
  
      // iterate over records to generate spans
      for (ConsumerRecord<?, ?> record : records) {
        assertThat(record.value()).isNull();
        assertThat(record.key()).isNull();
      }
  
      testing.waitAndAssertTraces(
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span ->
                      span.hasName(SHARED_TOPIC + " publish")
                          .hasKind(SpanKind.PRODUCER)
                          .hasNoParent()
                          .hasAttributesSatisfyingExactly(sendAttributes(null, null, false)),
                  span ->
                      span.hasName(SHARED_TOPIC + " process")
                          .hasKind(SpanKind.CONSUMER)
                          .hasParent(trace.getSpan(0))
                          .hasAttributesSatisfyingExactly(
                              processAttributes(null, null, false, false))));
    }
  
    @Test
    void testRecordsWithTopicPartitionKafkaConsume()
        throws ExecutionException, InterruptedException, TimeoutException {
      String greeting = "Hello from MockConsumer!";
      producer
          .send(new ProducerRecord<>(SHARED_TOPIC, partition, null, greeting))
          .get(5, TimeUnit.SECONDS);
  
      testing.waitForTraces(1);
  
      awaitUntilConsumerIsReady();
      ConsumerRecords<?, ?> consumerRecords = poll(Duration.ofSeconds(5));
      List<? extends ConsumerRecord<?, ?>> recordsInPartition =
          consumerRecords.records(KafkaClientBaseTest.topicPartition);
      assertThat(recordsInPartition.size()).isEqualTo(1);
  
      // iterate over records to generate spans
      for (ConsumerRecord<?, ?> record : recordsInPartition) {
        assertThat(record.value()).isEqualTo(greeting);
        assertThat(record.key()).isNull();
      }
  
      testing.waitAndAssertTraces(
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span ->
                      span.hasName(SHARED_TOPIC + " publish")
                          .hasKind(SpanKind.PRODUCER)
                          .hasNoParent()
                          .hasAttributesSatisfyingExactly(sendAttributes(null, greeting, false)),
                  span ->
                      span.hasName(SHARED_TOPIC + " process")
                          .hasKind(SpanKind.CONSUMER)
                          .hasParent(trace.getSpan(0))
                          .hasAttributesSatisfyingExactly(
                              processAttributes(null, greeting, false, false))));
    }
  }
  