content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements. See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License. You may obtain a copy of the License at
   *
   *    http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.kafka.common.test.junit;
  
  import org.apache.kafka.common.network.ListenerName;
  import org.apache.kafka.common.test.api.AutoStart;
  import org.apache.kafka.common.test.api.ClusterConfig;
  import org.apache.kafka.common.test.api.ClusterConfigProperty;
  import org.apache.kafka.common.test.api.ClusterFeature;
  import org.apache.kafka.common.test.api.ClusterTemplate;
  import org.apache.kafka.common.test.api.ClusterTest;
  import org.apache.kafka.common.test.api.ClusterTestDefaults;
  import org.apache.kafka.common.test.api.ClusterTests;
  import org.apache.kafka.common.test.api.DetectThreadLeak;
  import org.apache.kafka.common.test.api.Type;
  import org.apache.kafka.server.common.Feature;
  import org.apache.kafka.server.util.timer.SystemTimer;
  
  import org.junit.jupiter.api.extension.AfterEachCallback;
  import org.junit.jupiter.api.extension.BeforeEachCallback;
  import org.junit.jupiter.api.extension.ExtensionContext;
  import org.junit.jupiter.api.extension.ExtensionContext.Namespace;
  import org.junit.jupiter.api.extension.ExtensionContext.Store;
  import org.junit.jupiter.api.extension.TestTemplateInvocationContext;
  import org.junit.jupiter.api.extension.TestTemplateInvocationContextProvider;
  import org.junit.platform.commons.util.ReflectionUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.lang.reflect.Method;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Set;
  import java.util.function.Function;
  import java.util.stream.Collectors;
  import java.util.stream.IntStream;
  import java.util.stream.Stream;
  
  /**
   * This class is a custom JUnit extension that will generate some number of test invocations depending on the processing
   * of a few custom annotations. These annotations are placed on so-called test template methods. Template methods look
   * like normal JUnit test methods, but instead of being invoked directly, they are used as templates for generating
   * multiple test invocations.
   *
   * Test class that use this extension should use one of the following annotations on each template method:
   *
   * <ul>
   *     <li>{@link ClusterTest}, define a single cluster configuration</li>
   *     <li>{@link ClusterTests}, provide multiple instances of @ClusterTest</li>
   *     <li>{@link ClusterTemplate}, define a static method that generates cluster configurations</li>
   * </ul>
   *
   * Any combination of these annotations may be used on a given test template method. If no test invocations are
   * generated after processing the annotations, an error is thrown.
   *
   * Depending on which annotations are used, and what values are given, different {@link ClusterConfig} will be
   * generated. Each ClusterConfig is used to create an underlying Kafka cluster that is used for the actual test
   * invocation.
   *
   * For example:
   *
   * <pre>
   * class SomeIntegrationTest {
   *   &#64;ClusterTest(brokers = 1, controllers = 1, clusterType = ClusterType.Both)
   *   def someTest(): Unit = {
   *     assertTrue(condition)
   *   }
   * }
   * </pre>
   *
   * will generate two invocations of "someTest" (since ClusterType.Both was given). For each invocation, the test class
   * SomeIntegrationTest will be instantiated, lifecycle methods (before/after) will be run, and "someTest" will be invoked.
   *
   * A special system property "kafka.cluster.test.repeat" can be used to cause repeated invocation of the tests.
   *
   * For example:
   *
   * <pre>
   * ./gradlew -Pkafka.cluster.test.repeat=3 :core:test
   * </pre>
   *
   * will cause all ClusterTest-s in the :core module to be invoked three times.
   */
  public class ClusterTestExtensions implements TestTemplateInvocationContextProvider, BeforeEachCallback, AfterEachCallback {
      public static final String CLUSTER_TEST_REPEAT_SYSTEM_PROP = "kafka.cluster.test.repeat";
  
      private static final String METRICS_METER_TICK_THREAD_PREFIX = "metrics-meter-tick-thread";
      private static final String SCALA_THREAD_PREFIX = "scala-";
      private static final String FORK_JOIN_POOL_THREAD_PREFIX = "ForkJoinPool";
      private static final String JUNIT_THREAD_PREFIX = "junit-";
      private static final String ATTACH_LISTENER_THREAD_PREFIX = "Attach Listener";
      private static final String PROCESS_REAPER_THREAD_PREFIX = "process reaper";
      private static final String RMI_THREAD_PREFIX = "RMI";
      private static final String DETECT_THREAD_LEAK_KEY = "detectThreadLeak";
      private static final Set<String> SKIPPED_THREAD_PREFIX = Set.of(METRICS_METER_TICK_THREAD_PREFIX, SCALA_THREAD_PREFIX,
              FORK_JOIN_POOL_THREAD_PREFIX, JUNIT_THREAD_PREFIX, ATTACH_LISTENER_THREAD_PREFIX, PROCESS_REAPER_THREAD_PREFIX,
              RMI_THREAD_PREFIX, SystemTimer.SYSTEM_TIMER_THREAD_PREFIX);
      private static final Logger log = LoggerFactory.getLogger(ClusterTestExtensions.class);
  
      @Override
      public boolean supportsTestTemplate(ExtensionContext context) {
          return true;
      }
  
      /**
       * Indicates whether this provider may return zero test template invocation contexts.
       * Prior to JUnit 5.12, returning zero contexts was silently allowed. Starting from JUnit 5.12,
       * a PreconditionViolationException is thrown unless this method explicitly returns {@code true}.
       *
       * @param context the extension context for the test template method
       * @return {@code true} to allow this provider to return zero invocation contexts
       * @see <a href="https://github.com/junit-team/junit5/commit/89a46dfa10c6447ef010fbff7903bfcb3c18975a">JUnit 5.12 Breaking Change</a>
       */
      public boolean mayReturnZeroTestTemplateInvocationContexts(ExtensionContext context) {
          return true;
      }
  
      private boolean isClusterTest(ExtensionContext context) {
          Method method = context.getRequiredTestMethod();
          return method.getDeclaredAnnotation(ClusterTemplate.class) != null ||
              method.getDeclaredAnnotation(ClusterTest.class) != null ||
              method.getDeclaredAnnotation(ClusterTests.class) != null;
      }
  
      @Override
      public Stream<TestTemplateInvocationContext> provideTestTemplateInvocationContexts(ExtensionContext context) {
          ClusterTestDefaults defaults = getClusterTestDefaults(context.getRequiredTestClass());
          List<TestTemplateInvocationContext> generatedContexts = new ArrayList<>();
  
          // Process the @ClusterTemplate annotation
          ClusterTemplate clusterTemplateAnnot = context.getRequiredTestMethod().getDeclaredAnnotation(ClusterTemplate.class);
          if (clusterTemplateAnnot != null) {
              generatedContexts.addAll(processClusterTemplate(context, clusterTemplateAnnot));
          }
  
          // Process single @ClusterTest annotation
          ClusterTest clusterTestAnnot = context.getRequiredTestMethod().getDeclaredAnnotation(ClusterTest.class);
          if (clusterTestAnnot != null) {
              generatedContexts.addAll(processClusterTests(context, new ClusterTest[]{clusterTestAnnot}, defaults));
          }
  
          // Process multiple @ClusterTest annotation within @ClusterTests
          ClusterTests clusterTestsAnnot = context.getRequiredTestMethod().getDeclaredAnnotation(ClusterTests.class);
          if (clusterTestsAnnot != null) {
              generatedContexts.addAll(processClusterTests(context, clusterTestsAnnot.value(), defaults));
          }
  
          return generatedContexts.stream();
      }
  
      @Override
      public void beforeEach(ExtensionContext context) {
          if (isClusterTest(context)) {
              DetectThreadLeak detectThreadLeak = DetectThreadLeak.of(thread ->
                  SKIPPED_THREAD_PREFIX.stream().noneMatch(prefix -> thread.getName().startsWith(prefix)));
              getStore(context).put(DETECT_THREAD_LEAK_KEY, detectThreadLeak);
          }
      }
  
      @Override
      public void afterEach(ExtensionContext context) {
          if (isClusterTest(context)) {
              DetectThreadLeak detectThreadLeak = getStore(context).remove(DETECT_THREAD_LEAK_KEY, DetectThreadLeak.class);
              if (detectThreadLeak == null) {
                  return;
              }
              List<Thread> threads = detectThreadLeak.newThreads();
              if (!threads.isEmpty()) {
                  log.warn("Thread leak detected: {}",
                      threads.stream().map(Thread::getName).collect(Collectors.joining(", ")));
              }
              // TODO: Inkless Infinispan cache leaks threads, so we can't assert this yet
              // assertTrue(threads.isEmpty(), "Thread leak detected: " +
              //    threads.stream().map(Thread::getName).collect(Collectors.joining(", ")));
          }
      }
  
      private Store getStore(ExtensionContext context) {
          return context.getStore(Namespace.create(context.getUniqueId()));
      }
  
      private int getTestRepeatCount() {
          int count;
          try {
              String repeatCount = System.getProperty(CLUSTER_TEST_REPEAT_SYSTEM_PROP, "1");
              count = Integer.parseInt(repeatCount);
          } catch (NumberFormatException e) {
              count = 1;
          }
          return count;
      }
  
      private TestTemplateInvocationContext invocationContextForClusterType(
          Type type,
          String baseDisplayName,
          ClusterConfig config
      ) {
          return switch (type) {
              case KRAFT -> new RaftClusterInvocationContext(baseDisplayName, config, false);
              case CO_KRAFT -> new RaftClusterInvocationContext(baseDisplayName, config, true);
          };
      }
  
      List<TestTemplateInvocationContext> processClusterTemplate(ExtensionContext context, ClusterTemplate annot) {
          if (annot.value().trim().isEmpty()) {
              throw new IllegalStateException("ClusterTemplate value can't be empty string.");
          }
  
          String baseDisplayName = context.getRequiredTestMethod().getName();
          int repeatCount = getTestRepeatCount();
          List<TestTemplateInvocationContext> contexts = IntStream.range(0, repeatCount)
              .mapToObj(__ -> generateClusterConfiguration(context, annot.value()).stream())
              .flatMap(Function.identity())
              .flatMap(config -> config.clusterTypes().stream().map(type -> invocationContextForClusterType(type, baseDisplayName, config)))
              .collect(Collectors.toList());
  
          if (contexts.isEmpty()) {
              throw new IllegalStateException("ClusterConfig generator method should provide at least one config");
          }
  
          return contexts;
      }
  
      @SuppressWarnings("unchecked")
      private List<ClusterConfig> generateClusterConfiguration(
          ExtensionContext context,
          String generateClustersMethods
      ) {
          Object testInstance = context.getTestInstance().orElse(null);
          Method method = ReflectionUtils.getRequiredMethod(context.getRequiredTestClass(), generateClustersMethods);
          return (List<ClusterConfig>) ReflectionUtils.invokeMethod(method, testInstance);
      }
  
      private List<TestTemplateInvocationContext> processClusterTests(
          ExtensionContext context,
          ClusterTest[] clusterTests,
          ClusterTestDefaults defaults
      ) {
          int repeatCount = getTestRepeatCount();
          List<TestTemplateInvocationContext> ret = IntStream.range(0, repeatCount)
              .mapToObj(__ -> Arrays.stream(clusterTests))
              .flatMap(Function.identity())
              .flatMap(clusterTest -> processClusterTestInternal(context, clusterTest, defaults).stream())
              .collect(Collectors.toList());
  
          if (ret.isEmpty()) {
              throw new IllegalStateException("processClusterTests method should provide at least one config");
          }
  
          return ret;
      }
  
      private List<TestTemplateInvocationContext> processClusterTestInternal(
          ExtensionContext context,
          ClusterTest clusterTest,
          ClusterTestDefaults defaults
      ) {
          Type[] types = clusterTest.types().length == 0 ? defaults.types() : clusterTest.types();
          Map<String, String> serverProperties = Stream.concat(Arrays.stream(defaults.serverProperties()), Arrays.stream(clusterTest.serverProperties()))
              .filter(e -> e.id() == -1)
              .collect(Collectors.toMap(ClusterConfigProperty::key, ClusterConfigProperty::value, (a, b) -> b));
  
          Map<Integer, Map<String, String>> perServerProperties = Stream.concat(Arrays.stream(defaults.serverProperties()), Arrays.stream(clusterTest.serverProperties()))
              .filter(e -> e.id() != -1)
              .collect(Collectors.groupingBy(ClusterConfigProperty::id, Collectors.mapping(Function.identity(),
                  Collectors.toMap(ClusterConfigProperty::key, ClusterConfigProperty::value, (a, b) -> b))));
  
          Map<Feature, Short> features = Arrays.stream(clusterTest.features())
              .collect(Collectors.toMap(ClusterFeature::feature, ClusterFeature::version));
  
          ClusterConfig config = ClusterConfig.builder()
              .setTypes(Set.of(types))
              .setBrokers(clusterTest.brokers() == 0 ? defaults.brokers() : clusterTest.brokers())
              .setControllers(clusterTest.controllers() == 0 ? defaults.controllers() : clusterTest.controllers())
              .setDisksPerBroker(clusterTest.disksPerBroker() == 0 ? defaults.disksPerBroker() : clusterTest.disksPerBroker())
              .setAutoStart(clusterTest.autoStart() == AutoStart.DEFAULT ? defaults.autoStart() : clusterTest.autoStart() == AutoStart.YES)
              .setBrokerListenerName(ListenerName.normalised(clusterTest.brokerListener()))
              .setBrokerSecurityProtocol(clusterTest.brokerSecurityProtocol())
              .setControllerListenerName(ListenerName.normalised(clusterTest.controllerListener()))
              .setControllerSecurityProtocol(clusterTest.controllerSecurityProtocol())
              .setServerProperties(serverProperties)
              .setPerServerProperties(perServerProperties)
              .setMetadataVersion(clusterTest.metadataVersion())
              .setTags(List.of(clusterTest.tags()))
              .setFeatures(features)
              .build();
  
          return Arrays.stream(types)
              .map(type -> invocationContextForClusterType(type, context.getRequiredTestMethod().getName(), config))
              .collect(Collectors.toList());
      }
  
      private ClusterTestDefaults getClusterTestDefaults(Class<?> testClass) {
          return Optional.ofNullable(testClass.getDeclaredAnnotation(ClusterTestDefaults.class))
              .orElseGet(() -> EmptyClass.class.getDeclaredAnnotation(ClusterTestDefaults.class));
      }
  
      @ClusterTestDefaults
      private static final class EmptyClass {
          // Just used as a convenience to get default values from the annotation
      }
  }
  