Extend HttpRequestBuilder with additional setters

1 files changed, 20 insertions, 1 deletions

diff --git a/container-core/src/main/java/com/yahoo/container/jdisc/HttpRequestBuilder.java b/container-core/src/main/java/com/yahoo/container/jdisc/HttpRequestBuilder.java
index a2d792e6ae0..147f388e08c 100644
--- a/container-core/src/main/java/com/yahoo/container/jdisc/HttpRequestBuilder.java
+++ b/container-core/src/main/java/com/yahoo/container/jdisc/HttpRequestBuilder.java
@@ -4,6 +4,8 @@ package com.yahoo.container.jdisc;
 import com.yahoo.jdisc.http.HttpRequest.Method;
 
 import java.io.InputStream;
+import java.net.SocketAddress;
+import java.security.Principal;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -19,9 +21,13 @@ public class HttpRequestBuilder {
     private final String path;
     private final Map<String, List<String>> queryParameters = new TreeMap<>();
     private final Map<String, String> headers = new TreeMap<>();
+    private final Map<String, Object> attributes = new TreeMap<>();
     private String scheme;
     private String hostname;
     private InputStream content;
+    private Principal principal;
+    private SocketAddress socketAddress;
+    private int port = -1;
 
     private HttpRequestBuilder(Method method, String path) {
         this.method = method;
@@ -43,10 +49,20 @@ public class HttpRequestBuilder {
 
     public HttpRequestBuilder withHostname(String hostname) { this.hostname = hostname; return this; }
 
+    public HttpRequestBuilder withPrincipal(Principal p) { principal = p; return this; }
+
+    public HttpRequestBuilder withRemoteAddress(SocketAddress sa) { socketAddress = sa; return this; }
+
+    public HttpRequestBuilder withAttribute(String name, Object value) { attributes.put(name, value); return this; }
+
+    public HttpRequestBuilder withPort(int port) { this.port = port; return this; }
+
     public HttpRequest build() {
         String scheme = this.scheme != null ? this.scheme : "http";
         String hostname = this.hostname != null ? this.hostname : "localhost";
-        StringBuilder uriBuilder = new StringBuilder(scheme).append("://").append(hostname).append(path);
+        StringBuilder uriBuilder = new StringBuilder(scheme).append("://").append(hostname);
+        if (port > 0) uriBuilder.append(':').append(port);
+        uriBuilder.append(path);
         if (queryParameters.size() > 0) {
             uriBuilder.append('?');
             queryParameters.forEach((name, values) -> {
@@ -66,6 +82,9 @@ public class HttpRequestBuilder {
             request = HttpRequest.createTestRequest(uriBuilder.toString(), method);
         }
         headers.forEach((name, value) -> request.getJDiscRequest().headers().put(name, value));
+        if (principal != null) request.getJDiscRequest().setUserPrincipal(principal);
+        if (socketAddress != null) request.getJDiscRequest().setRemoteAddress(socketAddress);
+        request.getJDiscRequest().context().putAll(attributes);
         return request;
     }
 }