Added ability to load extensions via ServiceLoader, both from factories and directly

Author: Tom Akehurst

build.gradle

@@ -132,6 +132,8 @@ dependencies {
   }
   testImplementation "io.netty:netty-all:4.1.93.Final"
 
+  testImplementation files('test-extension/test-extension.jar')
+
   constraints {
     implementation "net.minidev:json-smart:2.4.11", {
       because 'Pinning this above the transitive version from json-path to get CVE fix'

src/main/java/com/github/tomakehurst/wiremock/core/WireMockConfiguration.java

@@ -411,7 +411,7 @@ public WireMockConfiguration extensions(Extension... extensionInstances) {
     return this;
   }
 
-  public WireMockConfiguration extensions(ExtensionFactory<?>... extensionFactories) {
+  public WireMockConfiguration extensions(ExtensionFactory... extensionFactories) {
     extensions.add(extensionFactories);
     return this;
   }

src/main/java/com/github/tomakehurst/wiremock/extension/ExtensionDeclarations.java

@@ -24,7 +24,7 @@ public class ExtensionDeclarations {
   private final List<String> classNames;
   private final List<Class<? extends Extension>> classes;
   private final Map<String, Extension> instances;
-  private final List<ExtensionFactory<? extends Extension>> factories;
+  private final List<ExtensionFactory> factories;
 
   public ExtensionDeclarations() {
     this.classNames = new ArrayList<>();
@@ -45,7 +45,7 @@ public void add(Class<? extends Extension>... classes) {
     this.classes.addAll(asList(classes));
   }
 
-  public void add(ExtensionFactory<? extends Extension>... factories) {
+  public void add(ExtensionFactory... factories) {
     this.factories.addAll(asList(factories));
   }
 
@@ -61,7 +61,7 @@ public Map<String, Extension> getInstances() {
     return instances;
   }
 
-  public List<ExtensionFactory<? extends Extension>> getFactories() {
+  public List<ExtensionFactory> getFactories() {
     return factories;
   }
 }

src/main/java/com/github/tomakehurst/wiremock/extension/ExtensionFactory.java

@@ -15,6 +15,8 @@
  */
 package com.github.tomakehurst.wiremock.extension;
 
-public interface ExtensionFactory<T extends Extension> {
-  T create(WireMockServices services);
+import java.util.List;
+
+public interface ExtensionFactory {
+  List<Extension> create(WireMockServices services);
 }

src/main/java/com/github/tomakehurst/wiremock/extension/Extensions.java

@@ -28,10 +28,9 @@
 import com.github.tomakehurst.wiremock.store.Stores;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Streams;
-import java.util.LinkedHashMap;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 import java.util.function.Function;
+import java.util.stream.Stream;
 
 public class Extensions implements WireMockServices {
 
@@ -80,14 +79,32 @@ public void load() {
             });
 
     loadedExtensions.putAll(extensionDeclarations.getInstances());
+
+    loadedExtensions.putAll(
+        loadExtensionsAsServices().collect(toMap(Extension::getName, Function.identity())));
+
+    final Stream<ExtensionFactory> allFactories =
+        Streams.concat(
+            extensionDeclarations.getFactories().stream(), loadExtensionFactoriesAsServices());
     loadedExtensions.putAll(
-        extensionDeclarations.getFactories().stream()
+        allFactories
             .map(factory -> factory.create(Extensions.this))
+            .flatMap(List::stream)
             .collect(toMap(Extension::getName, Function.identity())));
 
     configureTemplating();
   }
 
+  private Stream<Extension> loadExtensionsAsServices() {
+    final ServiceLoader<Extension> loader = ServiceLoader.load(Extension.class);
+    return loader.stream().map(ServiceLoader.Provider::get);
+  }
+
+  private Stream<ExtensionFactory> loadExtensionFactoriesAsServices() {
+    final ServiceLoader<ExtensionFactory> loader = ServiceLoader.load(ExtensionFactory.class);
+    return loader.stream().map(ServiceLoader.Provider::get);
+  }
+
   private void configureTemplating() {
     final Map<String, Helper<?>> helpers =
         ofType(TemplateHelperProviderExtension.class).values().stream()

src/main/java/com/github/tomakehurst/wiremock/stubbing/AbstractStubMappings.java

@@ -16,13 +16,15 @@
 package com.github.tomakehurst.wiremock.stubbing;
 
 import static com.github.tomakehurst.wiremock.common.LocalNotifier.notifier;
+import static com.github.tomakehurst.wiremock.common.Pair.pair;
 import static com.github.tomakehurst.wiremock.http.ResponseDefinition.copyOf;
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static java.util.stream.Collectors.toList;
 
 import com.github.tomakehurst.wiremock.admin.NotFoundException;
 import com.github.tomakehurst.wiremock.common.FileSource;
 import com.github.tomakehurst.wiremock.common.Json;
+import com.github.tomakehurst.wiremock.common.Pair;
 import com.github.tomakehurst.wiremock.extension.Parameters;
 import com.github.tomakehurst.wiremock.extension.ResponseDefinitionTransformer;
 import com.github.tomakehurst.wiremock.extension.ResponseDefinitionTransformerV2;
@@ -94,8 +96,11 @@ public ServeEvent serveFor(ServeEvent initialServeEvent) {
         initialServeEvent
             .withStubMapping(matchingMapping)
             .withResponseDefinition(responseDefinition);
-    responseDefinition =
+
+    final Pair<ServeEvent, ResponseDefinition> transformed =
         applyV2Transformations(serveEvent, ImmutableList.copyOf(v2transformers.values()));
+    serveEvent = transformed.a;
+    responseDefinition = transformed.b;
 
     return serveEvent.withResponseDefinition(copyOf(responseDefinition));
   }
@@ -122,13 +127,13 @@ private ResponseDefinition applyV1Transformations(
         request, newResponseDef, transformers.subList(1, transformers.size()));
   }
 
-  private ResponseDefinition applyV2Transformations(
+  private Pair<ServeEvent, ResponseDefinition> applyV2Transformations(
       ServeEvent serveEvent, List<ResponseDefinitionTransformerV2> transformers) {
 
     final ResponseDefinition responseDefinition = serveEvent.getResponseDefinition();
 
     if (transformers.isEmpty()) {
-      return responseDefinition;
+      return pair(serveEvent, responseDefinition);
     }
 
     ResponseDefinitionTransformerV2 transformer = transformers.get(0);

src/test/java/com/github/tomakehurst/wiremock/ExtensionFactoryTest.java

@@ -15,12 +15,15 @@
  */
 package com.github.tomakehurst.wiremock;
 
+import static com.github.tomakehurst.wiremock.client.WireMock.get;
+import static com.github.tomakehurst.wiremock.client.WireMock.noContent;
 import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.options;
 import static com.github.tomakehurst.wiremock.http.RequestMethod.GET;
 import static com.github.tomakehurst.wiremock.testsupport.TestFiles.defaultTestFilesRoot;
 import static net.javacrumbs.jsonunit.JsonMatchers.*;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.endsWith;
+import static org.hamcrest.Matchers.is;
 
 import com.github.tomakehurst.wiremock.admin.Router;
 import com.github.tomakehurst.wiremock.common.FileSource;
@@ -33,6 +36,7 @@
 import com.github.tomakehurst.wiremock.http.ResponseDefinition;
 import com.github.tomakehurst.wiremock.store.Stores;
 import com.github.tomakehurst.wiremock.testsupport.WireMockTestClient;
+import java.util.List;
 import java.util.Map;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
@@ -59,17 +63,14 @@ void injectsCoreServicesOnConstructionByFactory() {
             .templatingEnabled(false)
             .extensions(
                 services ->
-                    new MiscInfoApi(
-                        services.getAdmin(),
-                        services.getOptions(),
-                        services.getStores(),
-                        services.getFiles(),
-                        services.getExtensions())));
-
-    assertCorrectValuesReturnedFromApiCall();
-  }
+                    List.of(
+                        new MiscInfoApi(
+                            services.getAdmin(),
+                            services.getOptions(),
+                            services.getStores(),
+                            services.getFiles(),
+                            services.getExtensions()))));
 
-  private void assertCorrectValuesReturnedFromApiCall() {
     client.get("/something");
     client.get("/something");
 
@@ -79,7 +80,31 @@ private void assertCorrectValuesReturnedFromApiCall() {
     assertThat(content, jsonPartMatches("fileSourcePath", endsWith("test-file-root/__files")));
     assertThat(content, jsonPartEquals("requestCount", 2));
     assertThat(content, jsonPartEquals("stubCorsEnabled", true));
-    assertThat(content, jsonPartEquals("extensionCount", 1));
+    assertThat(
+        content, jsonPartEquals("extensionCount", 3)); // Includes the two service loaded extensions
+  }
+
+  @Test
+  void usesExtensionFactoryLoadedViaServiceLoader() {
+    initialiseWireMockServer(
+        options().dynamicPort().withRootDirectory(defaultTestFilesRoot()).templatingEnabled(false));
+
+    wm.stubFor(get("/transform-this").willReturn(noContent().withTransformers("loader-test")));
+
+    client.get("/just-count-this");
+
+    assertThat(client.get("/transform-this").content(), is("Request count 1"));
+  }
+
+  @Test
+  void usesExtensionInstanceLoadedViaServiceLoader() {
+    initialiseWireMockServer(
+        options().dynamicPort().withRootDirectory(defaultTestFilesRoot()).templatingEnabled(false));
+
+    wm.stubFor(
+        get("/transform-this").willReturn(noContent().withTransformers("instance-loader-test")));
+
+    assertThat(client.get("/transform-this").content(), is("Expected stuff"));
   }
 
   private void initialiseWireMockServer(WireMockConfiguration options) {

src/test/java/com/github/tomakehurst/wiremock/ResponseTransformerV2AcceptanceTest.java

@@ -30,6 +30,7 @@
 import com.github.tomakehurst.wiremock.http.Response;
 import com.github.tomakehurst.wiremock.stubbing.ServeEvent;
 import com.github.tomakehurst.wiremock.testsupport.WireMockTestClient;
+import java.util.List;
 import org.junit.jupiter.api.Test;
 
 public class ResponseTransformerV2AcceptanceTest {
@@ -74,7 +75,8 @@ public void globalTransformAppliedWithLocalParameters() {
 
   @Test
   public void filesRootIsCorrectlyPassedToTransformer() {
-    startWithExtensions(services -> new FilesUsingResponseTransformer(services.getFiles()));
+    startWithExtensions(
+        services -> List.of(new FilesUsingResponseTransformer(services.getFiles())));
 
     wm.stubFor(get(urlEqualTo("/response-transform-with-files")).willReturn(ok()));
 
@@ -89,7 +91,7 @@ private void startWithExtensions(Class<? extends Extension> extensionClasses) {
     client = new WireMockTestClient(wm.port());
   }
 
-  private void startWithExtensions(ExtensionFactory<?>... extensionFactories) {
+  private void startWithExtensions(ExtensionFactory... extensionFactories) {
     wm = new WireMockServer(wireMockConfig().dynamicPort().extensions(extensionFactories));
     wm.start();
     client = new WireMockTestClient(wm.port());

src/test/java/com/github/tomakehurst/wiremock/extension/responsetemplating/helpers/HandlebarsJsonPathHelperTest.java

@@ -35,6 +35,7 @@
 import com.google.common.collect.ImmutableMap;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -274,21 +275,23 @@ public void rendersAMeaningfulErrorWhenJsonPathIsNull() {
   @Test
   public void extractsValueFromAMap() {
     ResponseTemplateTransformer transformer =
-        buildExtension(
-            new MockWireMockServices(),
-            services ->
-                new ResponseTemplateTransformer(
-                    services.getTemplateEngine(), true, services.getFiles()) {
-                  @Override
-                  protected Map<String, Object> addExtraModelElements(
-                      Request request,
-                      ResponseDefinition responseDefinition,
-                      FileSource files,
-                      Parameters parameters) {
-                    return ImmutableMap.<String, Object>of(
-                        "mapData", ImmutableMap.of("things", "abc"));
-                  }
-                });
+        (ResponseTemplateTransformer)
+            buildExtension(
+                new MockWireMockServices(),
+                services ->
+                    List.of(
+                        new ResponseTemplateTransformer(
+                            services.getTemplateEngine(), true, services.getFiles()) {
+                          @Override
+                          protected Map<String, Object> addExtraModelElements(
+                              Request request,
+                              ResponseDefinition responseDefinition,
+                              FileSource files,
+                              Parameters parameters) {
+                            return ImmutableMap.<String, Object>of(
+                                "mapData", ImmutableMap.of("things", "abc"));
+                          }
+                        }));
 
     final ResponseDefinition responseDefinition =
         transform(

src/test/java/com/github/tomakehurst/wiremock/testsupport/ExtensionFactoryUtils.java

@@ -22,44 +22,51 @@
 import com.github.tomakehurst.wiremock.extension.ExtensionFactory;
 import com.github.tomakehurst.wiremock.extension.responsetemplating.ResponseTemplateTransformer;
 import com.github.tomakehurst.wiremock.extension.responsetemplating.ResponseTemplateTransformerTest;
+import java.util.List;
 import java.util.Map;
 
 public class ExtensionFactoryUtils {
 
   public static ResponseTemplateTransformer buildTemplateTransformer(boolean global) {
-    return buildExtension(
-        new MockWireMockServices(),
-        services ->
-            new ResponseTemplateTransformer(
-                services.getTemplateEngine(), global, services.getFiles()));
+    return (ResponseTemplateTransformer)
+        buildExtension(
+            new MockWireMockServices(),
+            services ->
+                List.of(
+                    new ResponseTemplateTransformer(
+                        services.getTemplateEngine(), global, services.getFiles())));
   }
 
   public static ResponseTemplateTransformer buildTemplateTransformer(
       boolean global, String helperName, Helper<?> helper) {
-    return buildExtension(
-        new MockWireMockServices().setHelpers(Map.of(helperName, helper)),
-        services ->
-            new ResponseTemplateTransformer(
-                services.getTemplateEngine(), global, services.getFiles()));
+    return (ResponseTemplateTransformer)
+        buildExtension(
+            new MockWireMockServices().setHelpers(Map.of(helperName, helper)),
+            services ->
+                List.of(
+                    new ResponseTemplateTransformer(
+                        services.getTemplateEngine(), global, services.getFiles())));
   }
 
   public static ResponseTemplateTransformer buildTemplateTransformer(Long maxCacheEntries) {
-    return buildExtension(
-        new MockWireMockServices().setMaxCacheEntries(maxCacheEntries),
-        services ->
-            new ResponseTemplateTransformer(
-                services.getTemplateEngine(), false, services.getFiles()));
+    return (ResponseTemplateTransformer)
+        buildExtension(
+            new MockWireMockServices().setMaxCacheEntries(maxCacheEntries),
+            services ->
+                List.of(
+                    new ResponseTemplateTransformer(
+                        services.getTemplateEngine(), false, services.getFiles())));
   }
 
-  public static <T extends Extension> T buildExtension(
-      MockWireMockServices wireMockServices, ExtensionFactory<T> factory) {
+  public static Extension buildExtension(
+      MockWireMockServices wireMockServices, ExtensionFactory factory) {
     FileSource fileSource =
         new ClasspathFileSource(
             ResponseTemplateTransformerTest.class
                 .getClassLoader()
                 .getResource("templates")
                 .getPath());
     wireMockServices.setFileSource(fileSource);
-    return factory.create(wireMockServices);
+    return factory.create(wireMockServices).stream().findFirst().get();
   }
 }