1 files changed, 34 insertions, 11 deletions
diff --git a/vespa-http-client/src/main/java/com/yahoo/vespa/http/client/core/communication/DryRunGatewayConnection.java b/vespa-http-client/src/main/java/com/yahoo/vespa/http/client/core/communication/DryRunGatewayConnection.java
index 129fc000271..623ea543ffb 100644
--- a/vespa-http-client/src/main/java/com/yahoo/vespa/http/client/core/communication/DryRunGatewayConnection.java
+++ b/vespa-http-client/src/main/java/com/yahoo/vespa/http/client/core/communication/DryRunGatewayConnection.java
@@ -5,8 +5,10 @@ import com.yahoo.vespa.http.client.config.Endpoint;
 import com.yahoo.vespa.http.client.core.Document;
 import com.yahoo.vespa.http.client.core.ErrorCode;
 import com.yahoo.vespa.http.client.core.OperationStatus;
+import com.yahoo.vespa.http.client.core.ServerResponseException;
 
 import java.io.ByteArrayInputStream;
+import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
 import java.time.Clock;
@@ -29,7 +31,13 @@ public class DryRunGatewayConnection implements GatewayConnection {
 
     /** Set to true to hold off responding with a result to any incoming operations until this is set false */
     private boolean hold = false;
-    private List<Document> held = new ArrayList<>();
+    private final List<Document> held = new ArrayList<>();
+
+    /** If this is set, handshake operations will throw this exception */
+    private ServerResponseException throwThisOnHandshake = null;
+
+    /** If this is set, all write operations will throw this exception */
+    private IOException throwThisOnWrite = null;
 
     public DryRunGatewayConnection(Endpoint endpoint, Clock clock) {
         this.endpoint = endpoint;
@@ -37,27 +45,27 @@ public class DryRunGatewayConnection implements GatewayConnection {
     }
 
     @Override
-    public InputStream write(List<Document> docs) {
-        StringBuilder result = new StringBuilder();
+    public InputStream write(List<Document> docs) throws IOException {
+        if (throwThisOnWrite != null)
+            throw throwThisOnWrite;
+
         if (hold) {
             held.addAll(docs);
+            return new ByteArrayInputStream("".getBytes(StandardCharsets.UTF_8));
         }
         else {
+            StringBuilder result = new StringBuilder();
             for (Document doc : held)
                 result.append(okResponse(doc).render());
             held.clear();
             for (Document doc : docs)
                 result.append(okResponse(doc).render());
+            return new ByteArrayInputStream(result.toString().getBytes(StandardCharsets.UTF_8));
         }
-        return new ByteArrayInputStream(result.toString().getBytes(StandardCharsets.UTF_8));
-    }
-
-    public void hold(boolean hold) {
-        this.hold = hold;
     }
 
     @Override
-    public InputStream poll() {
+    public InputStream poll() throws IOException {
         lastPollTime = clock.instant();
         return write(new ArrayList<>());
     }
@@ -66,7 +74,7 @@ public class DryRunGatewayConnection implements GatewayConnection {
     public Instant lastPollTime() { return lastPollTime; }
 
     @Override
-    public InputStream drain() {
+    public InputStream drain() throws IOException {
         return write(new ArrayList<>());
     }
 
@@ -85,14 +93,29 @@ public class DryRunGatewayConnection implements GatewayConnection {
     }
 
     @Override
-    public void handshake() { }
+    public void handshake() throws ServerResponseException {
+        if (throwThisOnHandshake != null)
+            throw throwThisOnHandshake;
+    }
 
     @Override
     public void close() { }
 
+    public void hold(boolean hold) {
+        this.hold = hold;
+    }
+
     /** Returns the document currently held in this */
     public List<Document> held() { return Collections.unmodifiableList(held); }
 
+    public void throwOnWrite(IOException throwThisOnWrite) {
+        this.throwThisOnWrite = throwThisOnWrite;
+    }
+
+    public void throwOnHandshake(ServerResponseException throwThisOnHandshake) {
+        this.throwThisOnHandshake = throwThisOnHandshake;
+    }
+
     private OperationStatus okResponse(Document document) {
         return new OperationStatus("ok", document.getOperationId(), ErrorCode.OK, false, "");
     }