@Test
  public void testRequestDropping() {
    AtomicInteger nextCorrelationId = new AtomicInteger(1);
    Function<Integer, List<RequestInfo>> requestGen = numRequests -> IntStream.range(0, numRequests)
        .mapToObj(
            i -> new RequestInfo(sslHost, sslPort, new MockSend(nextCorrelationId.getAndIncrement()), replicaOnSslNode,
                null))
        .collect(Collectors.toList());
    List<ResponseInfo> responseInfoList;

    // Drop requests while the requests are waiting for a connection.
    // First poll will require connections to be created, so no responses will be returned.
    responseInfoList = networkClient.sendAndPoll(requestGen.apply(3), Collections.emptySet(), POLL_TIMEOUT_MS);
    Assert.assertEquals("No responses expected in first poll.", 0, responseInfoList.size());
    // Drop requests on the second poll. The requests should be removed from the pending request list as a result.
    responseInfoList =
        networkClient.sendAndPoll(Collections.emptyList(), new HashSet<>(Arrays.asList(2, 3)), POLL_TIMEOUT_MS);
    Assert.assertEquals("Should receive only as many responses as there were requests", 3, responseInfoList.size());
    for (ResponseInfo responseInfo : responseInfoList) {
      MockSend send = (MockSend) responseInfo.getRequestInfo().getRequest();
      if (send.getCorrelationId() == 1) {
        NetworkClientErrorCode error = responseInfo.getError();
        ByteBuf response = responseInfo.content();
        Assert.assertNull("Should not have encountered an error", error);
        Assert.assertNotNull("Should receive a valid response", response);
        int correlationIdInRequest = send.getCorrelationId();
        int correlationIdInResponse = response.readInt();
        Assert.assertEquals("Received response for the wrong request", correlationIdInRequest, correlationIdInResponse);
      } else {
        Assert.assertEquals("Expected connection unavailable on dropped request",
            NetworkClientErrorCode.ConnectionUnavailable, responseInfo.getError());
        Assert.assertNull("Should not receive a response", responseInfo.content());
      }
    }
    responseInfoList.forEach(ResponseInfo::release);

    // Test dropping of requests while the requests are in flight.
    // Set the selector to idle mode to prevent responses from coming back (even though connections are available at
    // this moment in time).
    selector.setState(MockSelectorState.IdlePoll);
    responseInfoList = networkClient.sendAndPoll(requestGen.apply(3), Collections.emptySet(), POLL_TIMEOUT_MS);
    Assert.assertEquals("No responses expected in idle poll.", 0, responseInfoList.size());
    // Set the selector back to normal mode and drop a request. It should be dropped by closing the connection.
    selector.setState(MockSelectorState.Good);
    responseInfoList = networkClient.sendAndPoll(Collections.emptyList(), Collections.singleton(4), POLL_TIMEOUT_MS);
    Assert.assertEquals("Should receive only as many responses as there were requests", 3, responseInfoList.size());
    for (ResponseInfo responseInfo : responseInfoList) {
      MockSend send = (MockSend) responseInfo.getRequestInfo().getRequest();
      if (send.getCorrelationId() != 4) {
        NetworkClientErrorCode error = responseInfo.getError();
        ByteBuf response = responseInfo.content();
        Assert.assertNull("Should not have encountered an error", error);
        Assert.assertNotNull("Should receive a valid response", response);
        int correlationIdInRequest = send.getCorrelationId();
        int correlationIdInResponse = response.readInt();
        Assert.assertEquals("Received response for the wrong request", correlationIdInRequest, correlationIdInResponse);
        responseInfo.release();
      } else {
        Assert.assertEquals("Expected network error (from closed connection for dropped request)",
            NetworkClientErrorCode.NetworkError, responseInfo.getError());
        Assert.assertNull("Should not receive a response", responseInfo.content());
        responseInfo.release();
      }
    }

    // Dropping a request that is not currently pending or in flight should be a no-op.
    responseInfoList = networkClient.sendAndPoll(Collections.emptyList(), Collections.singleton(1), POLL_TIMEOUT_MS);
    Assert.assertEquals("No more responses expected.", 0, responseInfoList.size());
    responseInfoList.forEach(ResponseInfo::release);
  }
Test dropping of requests by closing
ambry-network/src/test/java/com/github/ambry/network/SocketNetworkClientTest.java