Simplify Extentions api - Remove overhead with id

src/main/java/org/wiremock/integrations/testcontainers/WireMockContainer.java

@@ -23,12 +23,14 @@
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -64,7 +66,8 @@ public class WireMockContainer extends GenericContainer<WireMockContainer> {
 
     private final Map<String, Stub> mappingStubs = new HashMap<>();
     private final Map<String, MountableFile> mappingFiles = new HashMap<>();
-    private final Map<String, Extension> extensions = new HashMap<>();
+    private final Set<String> extensionClassNames = new LinkedHashSet<>();
+    private final Set<File> extensionJars = new HashSet<>();
 
     public WireMockContainer() {
         this(DEFAULT_TAG);
@@ -137,51 +140,59 @@ public WireMockContainer withFileFromResource(String name, Class<?> resource, St
 
     /**
      * Add extension that will be loaded from the specified JAR file.
-     * @param id Unique ID of the extension, for logging purposes
+     * @param classNames Class names of the extension to be included
+     * @param jar JAR to be included into the container
+     * @return this instance
+     */
+    public WireMockContainer withExtension(Collection<String> classNames, File jar) {
+        return withExtension(classNames, Collections.singleton(jar));
+    }
+
+    /**
+     * Add extension that will be loaded from the specified JAR file.
      * @param classNames Class names of the extension to be included
      * @param jars JARs to be included into the container
      * @return this instance
      */
-    public WireMockContainer withExtension(String id, Collection<String> classNames, Collection<File> jars) {
-        final Extension extension = new Extension(id);
-        extension.extensionClassNames.addAll(classNames);
-        extension.jars.addAll(jars);
-        extensions.put(id, extension);
+    public WireMockContainer withExtension(Collection<String> classNames, Collection<File> jars) {
+        extensionClassNames.addAll(classNames);
+        extensionJars.addAll(jars);
         return this;
     }
 
     /**
      * Add extension that will be loaded from the specified directory with JAR files.
-     * @param id Unique ID of the extension, for logging purposes
      * @param classNames Class names of the extension to be included
-     * @param jarDirectory Directory that stores all JARs
+     * @param jarsDirectory Directory that stores all JARs
      * @return this instance
      */
-    public WireMockContainer withExtension(String id, Collection<String> classNames, File jarDirectory) {
-        final List<File> jarsInTheDirectory;
-        try (Stream<Path> walk = Files.walk(jarDirectory.toPath())) {
-            jarsInTheDirectory = walk
+    public WireMockContainer withExtension(Collection<String> classNames, Path jarsDirectory) {
+        if (!Files.isDirectory(jarsDirectory)) {
+            throw new IllegalArgumentException("Path must refers to directory " + jarsDirectory);
+        }
+        try (Stream<Path> walk = Files.walk(jarsDirectory)) {
+
+            final List<File> jarsInTheDirectory = walk
                     .filter(p -> !Files.isDirectory(p))
                     .map(Path::toFile)
                     .filter(f -> f.toString().endsWith(".jar"))
                     .collect(Collectors.toList());
+            return withExtension(classNames, jarsInTheDirectory);
+
         } catch (IOException e) {
-            throw new IllegalArgumentException("Cannot list JARs in the directory " + jarDirectory, e);
+            throw new IllegalArgumentException("Cannot list JARs in the directory " + jarsDirectory, e);
         }
-
-        return withExtension(id, classNames, jarsInTheDirectory);
     }
 
     /**
      * Add extension that will be loaded from the classpath.
      * This method can be used if the extension is a part of the WireMock bundle,
-     * or a Jar is already added via {@link #withExtension(String, Collection, Collection)}}
-     * @param id Unique ID of the extension, for logging purposes
+     * or a Jar is already added via {@link #withExtension(Collection, Collection)}}
      * @param className Class name of the extension
      * @return this instance
      */
-    public WireMockContainer withExtension(String id, String className) {
-        return withExtension(id, Collections.singleton(className), Collections.emptyList());
+    public WireMockContainer withExtension(String className) {
+        return withExtension(Collections.singleton(className), Collections.emptyList());
     }
 
     public String getEndpoint() {
@@ -209,13 +220,8 @@ protected void configure() {
             withCopyToContainer(mount.getValue(), FILES_DIR + mount.getKey());
         }
 
-        final ArrayList<String> extensionClassNames = new ArrayList<>();
-        for (Map.Entry<String, Extension> entry : extensions.entrySet()) {
-            final Extension ext = entry.getValue();
-            extensionClassNames.addAll(ext.extensionClassNames);
-            for (File jar : ext.jars) {
-                withCopyToContainer(MountableFile.forHostPath(jar.toPath()), EXTENSIONS_DIR + jar.getName());
-            }
+        for (File jar : extensionJars) {
+            withCopyToContainer(MountableFile.forHostPath(jar.toPath()), EXTENSIONS_DIR + jar.getName());
         }
         if (!extensionClassNames.isEmpty()) {
             wireMockArgs.append(" --extensions ");
@@ -236,14 +242,4 @@ public Stub (String name, String json) {
         }
     }
 
-    private static final class Extension {
-        final String id;
-        final List<File> jars = new ArrayList<>();
-        final List<String> extensionClassNames = new ArrayList<>();
-
-        public Extension(String id) {
-            this.id = id;
-        }
-    }
-
 }

src/test/java/org/wiremock/integrations/testcontainers/WireMockContainerExtensionTest.java

@@ -43,9 +43,9 @@ public class WireMockContainerExtensionTest {
     public WireMockContainer wiremockServer = new WireMockContainer("2.35.0")
             .withStartupTimeout(Duration.ofSeconds(60))
             .withMapping("json-body-transformer", WireMockContainerExtensionTest.class, "json-body-transformer.json")
-            .withExtension("JSON Body Transformer",
+            .withExtension(
                     Collections.singleton("com.ninecookies.wiremock.extensions.JsonBodyTransformer"),
-                    Collections.singleton(Paths.get("target", "test-wiremock-extension", "wiremock-extensions-0.4.1-jar-with-dependencies.jar").toFile()));
+                    Paths.get("target", "test-wiremock-extension", "wiremock-extensions-0.4.1-jar-with-dependencies.jar").toFile());
 
     @Before
     public void before() {

src/test/java/org/wiremock/integrations/testcontainers/WireMockContainerExtensionsCombinationTest.java

@@ -27,7 +27,7 @@
 import java.net.http.HttpResponse;
 import java.nio.file.Paths;
 import java.time.Duration;
-import java.util.Collections;
+import java.util.List;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -42,12 +42,9 @@ public class WireMockContainerExtensionsCombinationTest {
     @Rule
     public WireMockContainer wiremockServer = new WireMockContainer("2.35.0")
             .withMapping("json-body-transformer", WireMockContainerExtensionsCombinationTest.class, "json-body-transformer.json")
-            .withExtension("Webhook",
-                    Collections.singleton("org.wiremock.webhooks.Webhooks"),
-                    Collections.singleton(Paths.get("target", "test-wiremock-extension", "wiremock-webhooks-extension-2.35.0.jar").toFile()))
-            .withExtension("JSON Body Transformer",
-                    Collections.singleton("com.ninecookies.wiremock.extensions.JsonBodyTransformer"),
-                    Collections.singleton(Paths.get("target", "test-wiremock-extension", "wiremock-extensions-0.4.1-jar-with-dependencies.jar").toFile()));
+            .withExtension(
+                    List.of("org.wiremock.webhooks.Webhooks", "com.ninecookies.wiremock.extensions.JsonBodyTransformer"),
+                    Paths.get("target", "test-wiremock-extension"));
 
     @Before
     public void before() {