Added ability for extension factories to be declared as classes

Author: tomakehurst

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

@@ -427,6 +427,10 @@ public WireMockConfiguration extensions(Extension... extensionInstances) {
     return this;
   }
 
+  public WireMockConfiguration extensionFactories(ExtensionFactory... extensionFactories) {
+    return extensions(extensionFactories);
+  }
+
   public WireMockConfiguration extensions(ExtensionFactory... extensionFactories) {
     extensions.add(extensionFactories);
     return this;
@@ -437,6 +441,12 @@ public WireMockConfiguration extensions(Class<? extends Extension>... classes) {
     return this;
   }
 
+  public WireMockConfiguration extensionFactories(
+      Class<? extends ExtensionFactory>... factoryClasses) {
+    extensions.addFactories(factoryClasses);
+    return this;
+  }
+
   public WireMockConfiguration extensionScanningEnabled(boolean enabled) {
     this.extensionScanningEnabled = enabled;
     return this;

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

@@ -26,6 +26,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<Class<? extends ExtensionFactory>> factoryClasses;
   private final List<ExtensionFactory> factories;
   private static final String WEBHOOK_MESSAGE =
       "Passing webhooks in extensions is no longer required and"
@@ -35,6 +36,7 @@ public ExtensionDeclarations() {
     this.classNames = new ArrayList<>();
     this.classes = new ArrayList<>();
     this.instances = new LinkedHashMap<>();
+    this.factoryClasses = new ArrayList<>();
     this.factories = new ArrayList<>();
   }
 
@@ -58,6 +60,10 @@ public void add(ExtensionFactory... factories) {
     this.factories.addAll(asList(factories));
   }
 
+  public void addFactories(Class<? extends ExtensionFactory>... factoryClasses) {
+    this.factoryClasses.addAll(asList(factoryClasses));
+  }
+
   public List<String> getClassNames() {
     return classNames;
   }
@@ -74,6 +80,10 @@ public List<ExtensionFactory> getFactories() {
     return factories;
   }
 
+  public List<Class<? extends ExtensionFactory>> getFactoryClasses() {
+    return factoryClasses;
+  }
+
   private boolean removeWebhook(String className) {
     if (className.equals(Webhooks.class.getName())) {
       System.out.println(WEBHOOK_MESSAGE);

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

@@ -20,6 +20,7 @@
 import static java.util.stream.Collectors.toMap;
 
 import com.github.jknack.handlebars.Helper;
+import com.github.tomakehurst.wiremock.common.Exceptions;
 import com.github.tomakehurst.wiremock.common.FileSource;
 import com.github.tomakehurst.wiremock.core.Admin;
 import com.github.tomakehurst.wiremock.core.Options;
@@ -93,11 +94,15 @@ public void load() {
           loadExtensionsAsServices().collect(toMap(Extension::getName, Function.identity())));
     }
 
+    final Stream<ExtensionFactory> declaredFactories =
+        Stream.concat(
+            extensionDeclarations.getFactories().stream(),
+            extensionDeclarations.getFactoryClasses().stream()
+                .map(Extensions::instantiateExtensionFactory));
     final Stream<ExtensionFactory> allFactories =
         options.isExtensionScanningEnabled()
-            ? Stream.concat(
-                extensionDeclarations.getFactories().stream(), loadExtensionFactoriesAsServices())
-            : extensionDeclarations.getFactories().stream();
+            ? Stream.concat(declaredFactories, loadExtensionFactoriesAsServices())
+            : declaredFactories;
 
     loadedExtensions.putAll(
         allFactories
@@ -255,4 +260,10 @@ public <T extends Extension> Map<String, T> ofType(Class<T> extensionType) {
                         (entry1, entry2) -> entry1,
                         LinkedHashMap::new)));
   }
+
+  private static ExtensionFactory instantiateExtensionFactory(
+      Class<? extends ExtensionFactory> factoryClass) {
+    return Exceptions.uncheck(
+        () -> factoryClass.getDeclaredConstructor().newInstance(), ExtensionFactory.class);
+  }
 }

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

@@ -15,8 +15,7 @@
  */
 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.client.WireMock.*;
 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;
@@ -26,15 +25,17 @@
 import static org.hamcrest.Matchers.is;
 
 import com.github.tomakehurst.wiremock.admin.Router;
+import com.github.tomakehurst.wiremock.client.ResponseDefinitionBuilder;
 import com.github.tomakehurst.wiremock.common.FileSource;
 import com.github.tomakehurst.wiremock.common.Strings;
 import com.github.tomakehurst.wiremock.core.Admin;
 import com.github.tomakehurst.wiremock.core.Options;
 import com.github.tomakehurst.wiremock.core.WireMockConfiguration;
-import com.github.tomakehurst.wiremock.extension.AdminApiExtension;
-import com.github.tomakehurst.wiremock.extension.Extensions;
+import com.github.tomakehurst.wiremock.extension.*;
 import com.github.tomakehurst.wiremock.http.ResponseDefinition;
 import com.github.tomakehurst.wiremock.store.Stores;
+import com.github.tomakehurst.wiremock.stubbing.ServeEvent;
+import com.github.tomakehurst.wiremock.testsupport.WireMockResponse;
 import com.github.tomakehurst.wiremock.testsupport.WireMockTestClient;
 import java.io.File;
 import java.util.List;
@@ -132,6 +133,19 @@ void usesExtensionInstanceLoadedViaServiceLoader() {
     assertThat(client.get("/transform-this").content(), is("Expected stuff"));
   }
 
+  @Test
+  @SuppressWarnings("unchecked")
+  void usesExtensionFactoryLoadedAsClass() {
+    initialiseWireMockServer(
+        options().dynamicPort().extensionFactories(HeaderAddingExtensionFactory.class));
+
+    wm.stubFor(any(anyUrl()).willReturn(ok()));
+
+    WireMockResponse response = client.get("/test");
+
+    assertThat(response.firstHeader("added-header"), is("present"));
+  }
+
   private void initialiseWireMockServer(WireMockConfiguration options) {
     wm = new WireMockServer(options);
     wm.start();
@@ -181,4 +195,30 @@ public void contributeAdminApiRoutes(Router router) {
           });
     }
   }
+
+  public static class HeaderAddingExtensionFactory implements ExtensionFactory {
+
+    @Override
+    public List<Extension> create(WireMockServices services) {
+      return List.of(
+          new ResponseDefinitionTransformerV2() {
+            @Override
+            public ResponseDefinition transform(ServeEvent serveEvent) {
+              return ResponseDefinitionBuilder.like(serveEvent.getResponseDefinition())
+                  .withHeader("added-header", "present")
+                  .build();
+            }
+
+            @Override
+            public boolean applyGlobally() {
+              return true;
+            }
+
+            @Override
+            public String getName() {
+              return "add-header";
+            }
+          });
+    }
+  }
 }