Supplier for JDK Http client builder (#104)

agorapulse · May 22, 2024 · d2e8617 · d2e8617
1 parent 322f82d
commit d2e8617
Showing 1 changed file with 22 additions and 10 deletions.
diff --git a/libs/gru/src/main/java/com/agorapulse/gru/http/Http.java b/libs/gru/src/main/java/com/agorapulse/gru/http/Http.java
@@ -27,6 +27,8 @@
 import java.time.Duration;
 import java.time.temporal.ChronoUnit;
 import java.util.function.Consumer;
+import java.util.function.Function;
+import java.util.function.Supplier;
 
 import static java.lang.management.ManagementFactory.getRuntimeMXBean;
 
@@ -44,25 +46,39 @@ public static Http create(Class<?> unitTestClass) {
     }
 
     public static Http create(Object unitTest, Consumer<HttpClient.Builder> configuration) {
-        return new Http(unitTest, configuration);
+        return new Http(unitTest, b -> {
+            configuration.accept(b);
+            return b;
+        });
     }
 
     public static Http create(Class<?> unitTestClass, Consumer<HttpClient.Builder> configuration) {
-        return new Http(unitTestClass, configuration);
+        return new Http(unitTestClass,  b -> {
+            configuration.accept(b);
+            return b;
+        });
     }
 
-    private final Consumer<HttpClient.Builder> configuration;
+    public static Http create(Object unitTest, Supplier<HttpClient.Builder> configuration) {
+        return new Http(unitTest, b -> configuration.get());
+    }
+
+    public static Http create(Class<?> unitTestClass, Supplier<HttpClient.Builder> configuration) {
+        return new Http(unitTestClass, b -> configuration.get());
+    }
+
+    private final Function<HttpClient.Builder, HttpClient.Builder> configuration;
 
     private GruHttpRequest request;
     private GruHttpResponse response;
 
-    private Http(Object unitTest, Consumer<HttpClient.Builder> configuration) {
+    private Http(Object unitTest, Function<HttpClient.Builder, HttpClient.Builder> configuration) {
         super(unitTest);
         this.configuration = configuration;
         reset();
     }
 
-    private Http(Class<?> unitTestClass, Consumer<HttpClient.Builder> configuration) {
+    private Http(Class<?> unitTestClass, Function<HttpClient.Builder, HttpClient.Builder> configuration) {
         super(unitTestClass);
         this.configuration = configuration;
         reset();
@@ -91,11 +107,7 @@ public void reset() {
     @Override
     public GruContext run(Squad squad, GruContext context) {
         try {
-            HttpClient.Builder builder = HttpClient.newBuilder();
-
-            if (configuration != null) {
-                configuration.accept(builder);
-            }
+            HttpClient.Builder builder = configuration == null ? HttpClient.newBuilder() : configuration.apply(HttpClient.newBuilder());
 
             if (isDebugMode()) {
                 increaseTimeouts(builder);