Added support for specifying extension factories via the configuratio…

…n annotation

src/main/java/org/wiremock/spring/ConfigureWireMock.java

@@ -3,6 +3,7 @@
 import com.github.tomakehurst.wiremock.WireMockServer;
 import com.github.tomakehurst.wiremock.core.WireMockConfiguration;
 import com.github.tomakehurst.wiremock.extension.Extension;
+import com.github.tomakehurst.wiremock.extension.ExtensionFactory;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 
@@ -100,6 +101,13 @@
    */
   Class<? extends Extension>[] extensions() default {};
 
+  /**
+   * WireMock extensions to register in {@link WireMockServer}.
+   *
+   * @return the extensions
+   */
+  Class<? extends ExtensionFactory>[] extensionFactories() default {};
+
   /**
    * Customizes {@link WireMockConfiguration} used by {@link WireMockServer} instance. Customizers
    * are ordered by their natural order in this array. Each customizer must have no-arg constructor.

src/main/java/org/wiremock/spring/internal/WireMockServerCreator.java

@@ -3,7 +3,9 @@
 import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.options;
 
 import com.github.tomakehurst.wiremock.WireMockServer;
+import com.github.tomakehurst.wiremock.common.Exceptions;
 import com.github.tomakehurst.wiremock.core.WireMockConfiguration;
+import com.github.tomakehurst.wiremock.extension.ExtensionFactory;
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.io.File;
 import java.nio.file.Path;
@@ -74,6 +76,10 @@ public WireMockServer createWireMockServer(
                       });
             });
 
+    if (options.extensionFactories().length > 0) {
+      serverOptions.extensionFactories(options.extensionFactories());
+    }
+
     if (options.extensions().length > 0) {
       serverOptions.extensions(options.extensions());
     }
@@ -277,4 +283,10 @@ private void applyCustomizers(
       }
     }
   }
+
+  private static ExtensionFactory instantiateExtensionFactory(
+      Class<? extends ExtensionFactory> factoryClass) {
+    return Exceptions.uncheck(
+        () -> factoryClass.getDeclaredConstructor().newInstance(), ExtensionFactory.class);
+  }
 }

src/test/java/test/ExtensionFactoriesTest.java

@@ -0,0 +1,79 @@
+package test;
+
+import static com.github.tomakehurst.wiremock.client.WireMock.*;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import com.github.tomakehurst.wiremock.WireMockServer;
+import com.github.tomakehurst.wiremock.client.ResponseDefinitionBuilder;
+import com.github.tomakehurst.wiremock.extension.Extension;
+import com.github.tomakehurst.wiremock.extension.ExtensionFactory;
+import com.github.tomakehurst.wiremock.extension.ResponseDefinitionTransformerV2;
+import com.github.tomakehurst.wiremock.extension.WireMockServices;
+import com.github.tomakehurst.wiremock.http.ResponseDefinition;
+import com.github.tomakehurst.wiremock.stubbing.ServeEvent;
+import java.net.URI;
+import java.net.http.HttpClient;
+import java.net.http.HttpRequest;
+import java.net.http.HttpResponse;
+import java.util.List;
+import org.junit.jupiter.api.Test;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.wiremock.spring.ConfigureWireMock;
+import org.wiremock.spring.EnableWireMock;
+import org.wiremock.spring.InjectWireMock;
+
+@SpringBootTest(classes = ExtensionFactoriesTest.AppConfiguration.class)
+@EnableWireMock({
+  @ConfigureWireMock(
+      extensionFactories = {ExtensionFactoriesTest.HeaderAddingExtensionFactory.class})
+})
+class ExtensionFactoriesTest {
+
+  @InjectWireMock private WireMockServer wm;
+
+  HttpClient httpClient = HttpClient.newHttpClient();
+
+  @Test
+  void extensionIsLoaded() throws Exception {
+    wm.stubFor(any(anyUrl()).willReturn(ok()));
+
+    HttpResponse<String> response =
+        httpClient.send(
+            HttpRequest.newBuilder().GET().uri(URI.create(wm.baseUrl() + "/test")).build(),
+            HttpResponse.BodyHandlers.ofString());
+
+    assertThat(response.headers().firstValue("added-header").get()).isEqualTo("present");
+  }
+
+  @SpringBootApplication
+  static class AppConfiguration {}
+
+  public static class HeaderAddingExtensionFactory implements ExtensionFactory {
+
+    @Override
+    public List<Extension> create(WireMockServices services) {
+      return List.of(new HeaderAddingExtension());
+    }
+  }
+
+  public static class HeaderAddingExtension implements 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 "header-add";
+    }
+  }
+}

src/test/java/test/ExtensionsTest.java

@@ -0,0 +1,64 @@
+package test;
+
+import static com.github.tomakehurst.wiremock.client.WireMock.*;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import com.github.tomakehurst.wiremock.WireMockServer;
+import com.github.tomakehurst.wiremock.client.ResponseDefinitionBuilder;
+import com.github.tomakehurst.wiremock.extension.ResponseDefinitionTransformerV2;
+import com.github.tomakehurst.wiremock.http.ResponseDefinition;
+import com.github.tomakehurst.wiremock.stubbing.ServeEvent;
+import java.net.URI;
+import java.net.http.HttpClient;
+import java.net.http.HttpRequest;
+import java.net.http.HttpResponse;
+import org.junit.jupiter.api.Test;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.wiremock.spring.ConfigureWireMock;
+import org.wiremock.spring.EnableWireMock;
+import org.wiremock.spring.InjectWireMock;
+
+@SpringBootTest(classes = ExtensionsTest.AppConfiguration.class)
+@EnableWireMock({@ConfigureWireMock(extensions = {ExtensionsTest.HeaderAddingExtension.class})})
+class ExtensionsTest {
+
+  @InjectWireMock private WireMockServer wm;
+
+  HttpClient httpClient = HttpClient.newHttpClient();
+
+  @Test
+  void extensionIsLoaded() throws Exception {
+    wm.stubFor(any(anyUrl()).willReturn(ok()));
+
+    HttpResponse<String> response =
+        httpClient.send(
+            HttpRequest.newBuilder().GET().uri(URI.create(wm.baseUrl() + "/test")).build(),
+            HttpResponse.BodyHandlers.ofString());
+
+    assertThat(response.headers().firstValue("added-header").get()).isEqualTo("present");
+  }
+
+  @SpringBootApplication
+  static class AppConfiguration {}
+
+  public static class HeaderAddingExtension implements 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 "header-add";
+    }
+  }
+}