4107: Added integration test testWriteChannelWithConnectionPool

andrey-qlogic · andrey-qlogic · commit 5e4aef0e113b · 2018-12-01T14:33:31.000Z
diff --git a/google-cloud-clients/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITStorageTest.java b/google-cloud-clients/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITStorageTest.java
@@ -28,9 +28,12 @@
 import static org.junit.Assume.assumeFalse;
 import static org.junit.Assume.assumeTrue;
 
+import com.google.api.client.http.HttpTransport;
+import com.google.api.client.http.apache.ApacheHttpTransport;
 import com.google.api.client.util.DateTime;
 import com.google.api.gax.paging.Page;
 import com.google.auth.ServiceAccountSigner;
+import com.google.auth.http.HttpTransportFactory;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.kms.v1.CreateCryptoKeyRequest;
 import com.google.cloud.kms.v1.CreateKeyRingRequest;
@@ -45,7 +48,9 @@
 import com.google.cloud.Policy;
 import com.google.cloud.ReadChannel;
 import com.google.cloud.RestorableState;
+import com.google.cloud.TransportOptions;
 import com.google.cloud.WriteChannel;
+import com.google.cloud.http.HttpTransportOptions;
 import com.google.cloud.kms.v1.LocationName;
 import com.google.cloud.storage.Acl;
 import com.google.cloud.storage.Acl.Role;
@@ -113,6 +118,8 @@
 import io.grpc.StatusRuntimeException;
 import io.grpc.stub.MetadataUtils;
 
+import org.apache.http.impl.client.HttpClients;
+import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Test;
@@ -184,6 +191,15 @@ public static void afterClass() throws ExecutionException, InterruptedException
     }
   }
 
+  private static class CustomHttpTransportFactory implements HttpTransportFactory {
+    @Override
+    public HttpTransport create() {
+      PoolingHttpClientConnectionManager manager = new PoolingHttpClientConnectionManager();
+      manager.setMaxTotal(1);
+      return new ApacheHttpTransport(HttpClients.createMinimal(manager));
+    }
+  }
+
   private static void prepareKmsKeys() throws IOException {
     String projectId = remoteStorageHelper.getOptions().getProjectId();
     GoogleCredentials credentials = GoogleCredentials.getApplicationDefault();
@@ -1741,6 +1757,27 @@ public void testWriteChannelExistingBlob() throws IOException {
     assertTrue(storage.delete(BUCKET, blobName));
   }
 
+  @Test(timeout = 5000)
+  public void testWriteChannelWithConnectionPool() throws IOException {
+    TransportOptions transportOptions = HttpTransportOptions.newBuilder()
+        .setHttpTransportFactory(new CustomHttpTransportFactory()).build();
+    Storage storageWithPool =
+        StorageOptions.newBuilder().setTransportOptions(transportOptions).build().getService();
+    String blobName = "test-read-and-write-channels-blob";
+    BlobInfo blob = BlobInfo.newBuilder(BUCKET, blobName).build();
+    byte[] stringBytes;
+    try (WriteChannel writer = storageWithPool.writer(blob)) {
+      stringBytes = BLOB_STRING_CONTENT.getBytes(UTF_8);
+      writer.write(ByteBuffer.wrap(BLOB_BYTE_CONTENT));
+      writer.write(ByteBuffer.wrap(stringBytes));
+    }
+    try (WriteChannel writer = storageWithPool.writer(blob)) {
+      stringBytes = BLOB_STRING_CONTENT.getBytes(UTF_8);
+      writer.write(ByteBuffer.wrap(BLOB_BYTE_CONTENT));
+      writer.write(ByteBuffer.wrap(stringBytes));
+    }
+  }
+
   @Test
   public void testGetSignedUrl() throws IOException {
     if(storage.getOptions().getCredentials() != null) {
