KAFKA-16654:Refactor kafka.test.annotation.Type and ClusterTestExtensions

core/src/test/java/kafka/test/annotation/Type.java

@@ -22,30 +22,31 @@
 import kafka.test.junit.ZkClusterInvocationContext;
 import org.junit.jupiter.api.extension.TestTemplateInvocationContext;
 
-import java.util.function.Consumer;
+import java.util.List;
+import java.util.Collections;
 
 /**
  * The type of cluster config being requested. Used by {@link kafka.test.ClusterConfig} and the test annotations.
  */
 public enum Type {
     KRAFT {
         @Override
-        public void invocationContexts(String baseDisplayName, ClusterConfig config, Consumer<TestTemplateInvocationContext> invocationConsumer) {
-            invocationConsumer.accept(new RaftClusterInvocationContext(baseDisplayName, config, false));
+        public List<TestTemplateInvocationContext> invocationContexts(String baseDisplayName, ClusterConfig config) {
+            return Collections.singletonList(new RaftClusterInvocationContext(baseDisplayName, config, false));
         }
     },
     CO_KRAFT {
         @Override
-        public void invocationContexts(String baseDisplayName, ClusterConfig config, Consumer<TestTemplateInvocationContext> invocationConsumer) {
-            invocationConsumer.accept(new RaftClusterInvocationContext(baseDisplayName, config, true));
+        public List<TestTemplateInvocationContext> invocationContexts(String baseDisplayName, ClusterConfig config) {
+            return Collections.singletonList(new RaftClusterInvocationContext(baseDisplayName, config, true));
         }
     },
     ZK {
         @Override
-        public void invocationContexts(String baseDisplayName, ClusterConfig config, Consumer<TestTemplateInvocationContext> invocationConsumer) {
-            invocationConsumer.accept(new ZkClusterInvocationContext(baseDisplayName, config));
+        public List<TestTemplateInvocationContext> invocationContexts(String baseDisplayName, ClusterConfig config) {
+            return Collections.singletonList(new ZkClusterInvocationContext(baseDisplayName, config));
         }
     };
 
-    public abstract void invocationContexts(String baseDisplayName, ClusterConfig config, Consumer<TestTemplateInvocationContext> invocationConsumer);
+    public abstract List<TestTemplateInvocationContext> invocationContexts(String baseDisplayName, ClusterConfig config);
 }

core/src/test/java/kafka/test/junit/ClusterTestExtensions.java

@@ -38,7 +38,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -94,38 +93,31 @@ public Stream<TestTemplateInvocationContext> provideTestTemplateInvocationContex
         // Process the @ClusterTemplate annotation
         ClusterTemplate clusterTemplateAnnot = context.getRequiredTestMethod().getDeclaredAnnotation(ClusterTemplate.class);
         if (clusterTemplateAnnot != null) {
-            processClusterTemplate(context, clusterTemplateAnnot, generatedContexts::add);
-            if (generatedContexts.isEmpty()) {
-                throw new IllegalStateException("ClusterConfig generator method should provide at least one config");
-            }
+            generatedContexts.addAll(processClusterTemplate(context, clusterTemplateAnnot));
         }
 
         // Process single @ClusterTest annotation
         ClusterTest clusterTestAnnot = context.getRequiredTestMethod().getDeclaredAnnotation(ClusterTest.class);
         if (clusterTestAnnot != null) {
-            processClusterTest(context, clusterTestAnnot, defaults, generatedContexts::add);
+            generatedContexts.addAll(processClusterTest(context, clusterTestAnnot, defaults));
         }
 
         // Process multiple @ClusterTest annotation within @ClusterTests
         ClusterTests clusterTestsAnnot = context.getRequiredTestMethod().getDeclaredAnnotation(ClusterTests.class);
         if (clusterTestsAnnot != null) {
-            for (ClusterTest annot : clusterTestsAnnot.value()) {
-                processClusterTest(context, annot, defaults, generatedContexts::add);
-            }
-        }
-
-        if (generatedContexts.isEmpty()) {
-            throw new IllegalStateException("Please annotate test methods with @ClusterTemplate, @ClusterTest, or " +
-                    "@ClusterTests when using the ClusterTestExtensions provider");
+            generatedContexts.addAll(processClusterTests(context, clusterTestsAnnot, defaults));
         }
 
         return generatedContexts.stream();
     }
 
-    void processClusterTemplate(ExtensionContext context, ClusterTemplate annot,
-                                        Consumer<TestTemplateInvocationContext> testInvocations) {
+
+
+    List<TestTemplateInvocationContext> processClusterTemplate(ExtensionContext context, ClusterTemplate annot) {
         // If specified, call cluster config generated method (must be static)
         List<ClusterConfig> generatedClusterConfigs = new ArrayList<>();
+        List<TestTemplateInvocationContext> contexts = new ArrayList<>();
+
         if (annot.value().trim().isEmpty()) {
             throw new IllegalStateException("ClusterTemplate value can't be empty string.");
         }
@@ -134,9 +126,15 @@ void processClusterTemplate(ExtensionContext context, ClusterTemplate annot,
         String baseDisplayName = context.getRequiredTestMethod().getName();
         generatedClusterConfigs.forEach(config -> {
             for (Type type: config.clusterTypes()) {
-                type.invocationContexts(baseDisplayName, config, testInvocations);
+                contexts.addAll(type.invocationContexts(baseDisplayName, config));
             }
         });
+
+        if (contexts.isEmpty()) {
+            throw new IllegalStateException("ClusterConfig generator method should provide at least one config");
+        }
+
+        return contexts;
     }
 
     private void generateClusterConfigurations(ExtensionContext context, String generateClustersMethods, ClusterGenerator generator) {
@@ -145,8 +143,30 @@ private void generateClusterConfigurations(ExtensionContext context, String gene
         ReflectionUtils.invokeMethod(method, testInstance, generator);
     }
 
-    private void processClusterTest(ExtensionContext context, ClusterTest annot, ClusterTestDefaults defaults,
-                                    Consumer<TestTemplateInvocationContext> testInvocations) {
+    private List<TestTemplateInvocationContext> processClusterTests(ExtensionContext context, ClusterTests annots, ClusterTestDefaults defaults) {
+        List<TestTemplateInvocationContext> ret = new ArrayList<>();
+
+        for (ClusterTest annot : annots.value()) {
+            ret.addAll(processClusterTestInternal(context, annot, defaults));
+        }
+
+        if (ret.isEmpty()) {
+            throw new IllegalStateException("processClusterTests method should provide at least one config");
+        }
+
+        return ret;
+    }
+
+    private List<TestTemplateInvocationContext> processClusterTest(ExtensionContext context, ClusterTest annot, ClusterTestDefaults defaults) {
+        List<TestTemplateInvocationContext> ret = processClusterTestInternal(context, annot, defaults);
+
+        if (ret.isEmpty()) {
+            throw new IllegalStateException("processClusterTest method should provide at least one config");
+        }
+
+        return ret;
+    }
+    private List<TestTemplateInvocationContext> processClusterTestInternal(ExtensionContext context, ClusterTest annot, ClusterTestDefaults defaults) {
         Type[] types = annot.types().length == 0 ? defaults.types() : annot.types();
         Map<String, String> serverProperties = Stream.concat(Arrays.stream(defaults.serverProperties()), Arrays.stream(annot.serverProperties()))
                 .filter(e -> e.id() == -1)
@@ -170,9 +190,12 @@ private void processClusterTest(ExtensionContext context, ClusterTest annot, Clu
                 .setSecurityProtocol(annot.securityProtocol())
                 .setMetadataVersion(annot.metadataVersion())
                 .build();
+
+        List<TestTemplateInvocationContext> ret = new ArrayList<>();
         for (Type type : types) {
-            type.invocationContexts(context.getRequiredTestMethod().getName(), config, testInvocations);
+            ret.addAll(type.invocationContexts(context.getRequiredTestMethod().getName(), config));
         }
+        return ret;
     }
 
     private ClusterTestDefaults getClusterTestDefaults(Class<?> testClass) {

core/src/test/java/kafka/test/junit/ClusterTestExtensionsUnitTest.java

@@ -22,8 +22,6 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.extension.ExtensionContext;
-import org.junit.jupiter.api.extension.TestTemplateInvocationContext;
-import java.util.function.Consumer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -33,16 +31,16 @@ public class ClusterTestExtensionsUnitTest {
     void testProcessClusterTemplate() {
         ClusterTestExtensions ext = new ClusterTestExtensions();
         ExtensionContext context = mock(ExtensionContext.class);
-        Consumer<TestTemplateInvocationContext> testInvocations = mock(Consumer.class);
+
         ClusterTemplate annot = mock(ClusterTemplate.class);
         when(annot.value()).thenReturn("").thenReturn(" ");
 
         Assertions.assertThrows(IllegalStateException.class, () ->
-                ext.processClusterTemplate(context, annot, testInvocations)
+                ext.processClusterTemplate(context, annot)
         );
 
         Assertions.assertThrows(IllegalStateException.class, () ->
-                ext.processClusterTemplate(context, annot, testInvocations)
+                ext.processClusterTemplate(context, annot)
         );
     }
 }