diff --git a/test/framework/src/main/java/org/elasticsearch/test/transport/CapturingTransport.java b/test/framework/src/main/java/org/elasticsearch/test/transport/CapturingTransport.java index 81fc934ca6d7e..318c70c2933d8 100644 --- a/test/framework/src/main/java/org/elasticsearch/test/transport/CapturingTransport.java +++ b/test/framework/src/main/java/org/elasticsearch/test/transport/CapturingTransport.java @@ -46,6 +46,7 @@ import java.io.UncheckedIOException; import java.net.UnknownHostException; import java.util.ArrayList; +import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.List; @@ -94,9 +95,17 @@ public CapturedRequest[] capturedRequests() { * @return the captured requests */ public CapturedRequest[] getCapturedRequestsAndClear() { - CapturedRequest[] capturedRequests = capturedRequests(); - clear(); - return capturedRequests; + List requests = new ArrayList<>(capturedRequests.size()); + capturedRequests.drainTo(requests); + return requests.toArray(new CapturedRequest[0]); + } + + private Map> groupRequestsByTargetNode(Collection requests) { + Map> result = new HashMap<>(); + for (CapturedRequest request : requests) { + result.computeIfAbsent(request.node.getId(), node -> new ArrayList<>()).add(request); + } + return result; } /** @@ -104,16 +113,7 @@ public CapturedRequest[] getCapturedRequestsAndClear() { * Doesn't clear the captured request list. See {@link #clear()} */ public Map> capturedRequestsByTargetNode() { - Map> map = new HashMap<>(); - for (CapturedRequest request : capturedRequests) { - List nodeList = map.get(request.node.getId()); - if (nodeList == null) { - nodeList = new ArrayList<>(); - map.put(request.node.getId(), nodeList); - } - nodeList.add(request); - } - return map; + return groupRequestsByTargetNode(capturedRequests); } /** @@ -125,9 +125,9 @@ public Map> capturedRequestsByTargetNode() { * @return the captured requests grouped by target node */ public Map> getCapturedRequestsByTargetNodeAndClear() { - Map> map = capturedRequestsByTargetNode(); - clear(); - return map; + List requests = new ArrayList<>(capturedRequests.size()); + capturedRequests.drainTo(requests); + return groupRequestsByTargetNode(requests); } /** clears captured requests */