allow use of progressmonitor for byte bodies

Kong · Oct 26, 2022 · 0b35541 · 0b35541
1 parent 7c1562d
commit 0b35541
Show file tree

Hide file tree

Showing 5 changed files with 90 additions and 38 deletions.
diff --git a/unirest-bdd-tests/src/test/java/BehaviorTests/DownloadProgressTest.java b/unirest-bdd-tests/src/test/java/BehaviorTests/DownloadProgressTest.java
@@ -64,4 +64,24 @@ void canAddUploadProgressAsync() throws Exception {
         monitor.assertSpideyFileDownload("spidey.jpg");
     }
 
+
+    @Test
+    void canAddUploadProgressWithBytes() {
+        Unirest.get(MockServer.BINARYFILE)
+                .downloadMonitor(monitor)
+                .asBytes();
+
+        monitor.assertSpideyFileDownload("body");
+    }
+
+    @Test
+    void canAddUploadProgressWithBytesAsync() throws Exception {
+        Unirest.get(MockServer.BINARYFILE)
+                .downloadMonitor(monitor)
+                .asBytesAsync()
+                .get();
+
+        monitor.assertSpideyFileDownload("body");
+    }
+
 }
diff --git a/unirest/src/main/java/kong/unirest/BaseRequest.java b/unirest/src/main/java/kong/unirest/BaseRequest.java
@@ -226,17 +226,17 @@ public CompletableFuture<HttpResponse<String>> asStringAsync(Callback<String> ca
 
     @Override
     public HttpResponse<byte[]> asBytes() {
-        return request(ByteResponse::new, byte[].class);
+        return request(r -> new ByteResponse(r, downloadMonitor), byte[].class);
     }
 
     @Override
     public CompletableFuture<HttpResponse<byte[]>> asBytesAsync() {
-        return config.getAsyncClient().request(this, ByteResponse::new, new CompletableFuture<>(), byte[].class);
+        return config.getAsyncClient().request(this, r -> new ByteResponse(r, downloadMonitor), new CompletableFuture<>(), byte[].class);
     }
 
     @Override
     public CompletableFuture<HttpResponse<byte[]>> asBytesAsync(Callback<byte[]> callback) {
-        return config.getAsyncClient().request(this, ByteResponse::new, wrap(callback), byte[].class);
+        return config.getAsyncClient().request(this, r -> new ByteResponse(r, downloadMonitor), wrap(callback), byte[].class);
     }
 
     @Override

diff --git a/unirest/src/main/java/kong/unirest/ByteResponse.java b/unirest/src/main/java/kong/unirest/ByteResponse.java
@@ -25,12 +25,54 @@
 
 package kong.unirest;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+
 public class ByteResponse extends BaseResponse<byte[]> {
     private final byte[] body;
 
-    public ByteResponse(RawResponse r) {
+    public ByteResponse(RawResponse r, ProgressMonitor downloadMonitor) {
         super(r);
-        this.body = r.getContentAsBytes();
+        if(downloadMonitor == null) {
+            this.body = r.getContentAsBytes();
+        } else {
+            MonitoringInputStream ip = new MonitoringInputStream(r.getContent(), downloadMonitor, (String)null, r);
+            try {
+                body = getBytes(ip);
+            } catch (IOException e){
+                throw new UnirestException(e);
+            }
+        }
+    }
+
+    public static byte[] getBytes(InputStream is) throws IOException {
+        try {
+            int len;
+            int size = 1024;
+            byte[] buf;
+
+            if (is instanceof ByteArrayInputStream) {
+                size = is.available();
+                buf = new byte[size];
+                len = is.read(buf, 0, size);
+            } else {
+                ByteArrayOutputStream bos = new ByteArrayOutputStream();
+                buf = new byte[size];
+                while ((len = is.read(buf, 0, size)) != -1) {
+                    bos.write(buf, 0, len);
+                }
+                buf = bos.toByteArray();
+            }
+            return buf;
+        } finally {
+            is.close();
+        }
+    }
+
+    public static boolean isGzipped(String value) {
+        return "gzip".equalsIgnoreCase(value.toLowerCase().trim());
     }
 
     @Override

diff --git a/unirest/src/main/java/kong/unirest/MonitoringInputStream.java b/unirest/src/main/java/kong/unirest/MonitoringInputStream.java
@@ -28,6 +28,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.file.Path;
+import java.util.zip.GZIPInputStream;
 
 class MonitoringInputStream extends InputStream {
     private final InputStream content;
@@ -36,11 +37,27 @@ class MonitoringInputStream extends InputStream {
     private long byteCount = 0;
     private String fileName;
 
-    MonitoringInputStream(InputStream content, ProgressMonitor downloadMonitor, Path target, RawResponse contentSize) {
-        this.content = content;
+    MonitoringInputStream(InputStream content, ProgressMonitor downloadMonitor, Path target, RawResponse rawResponse) {
+        this(content, downloadMonitor, target.getFileName().toString(), rawResponse);
+    }
+
+    MonitoringInputStream(InputStream content, ProgressMonitor downloadMonitor, String fileName, RawResponse rawResponse) {
+        this.content = wrap(content, rawResponse);
         this.downloadMonitor = downloadMonitor;
-        this.fileName = target.getFileName().toString();
-        this.totalSize = getBodySize(contentSize);
+        this.fileName = fileName;
+        this.totalSize = getBodySize(rawResponse);
+    }
+
+    private InputStream wrap(InputStream is , RawResponse rawResponse) {
+        try {
+            if (is.available() > 0 && "gzip".equalsIgnoreCase(rawResponse.getContentType())) {
+                return new GZIPInputStream(is);
+            } else {
+                return is;
+            }
+        }catch (Exception e){
+            throw new UnirestException(e);
+        }
     }
 
     private Long getBodySize(RawResponse r) {

diff --git a/unirest/src/main/java/kong/unirest/apache/ApacheResponse.java b/unirest/src/main/java/kong/unirest/apache/ApacheResponse.java
@@ -81,10 +81,10 @@ public byte[] getContentAsBytes() {
         }
         try {
             InputStream is = getContent();
-            if (is.available() > 0 && isGzipped(getEncoding())) {
+            if (is.available() > 0 && ByteResponse.isGzipped(getEncoding())) {
               is = new GZIPInputStream(getContent());
             }
-            return getBytes(is);
+            return ByteResponse.getBytes(is);
         } catch (IOException e2) {
             throw new UnirestException(e2);
         } finally {
@@ -149,31 +149,4 @@ public String getEncoding() {
         return "";
     }
 
-    private static byte[] getBytes(InputStream is) throws IOException {
-        try {
-            int len;
-            int size = 1024;
-            byte[] buf;
-
-            if (is instanceof ByteArrayInputStream) {
-                size = is.available();
-                buf = new byte[size];
-                len = is.read(buf, 0, size);
-            } else {
-                ByteArrayOutputStream bos = new ByteArrayOutputStream();
-                buf = new byte[size];
-                while ((len = is.read(buf, 0, size)) != -1) {
-                    bos.write(buf, 0, len);
-                }
-                buf = bos.toByteArray();
-            }
-            return buf;
-        } finally {
-            is.close();
-        }
-    }
-
-    private static boolean isGzipped(String value) {
-        return "gzip".equalsIgnoreCase(value.toLowerCase().trim());
-    }
 }