From 35f62f9960e575583b62a84506044cf69b57b8ef Mon Sep 17 00:00:00 2001
From: Peter Clement <peter@budibase.com>
Date: Mon, 18 Mar 2024 12:27:23 +0000
Subject: [PATCH] tidy up

---
 .../src/integrations/tests/rest.spec.ts       | 14 +++++++++-----
 packages/server/src/integrations/utils.ts     | 19 ++++++++++++++++---
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/packages/server/src/integrations/tests/rest.spec.ts b/packages/server/src/integrations/tests/rest.spec.ts
index 54ca87815f..09ae345337 100644
--- a/packages/server/src/integrations/tests/rest.spec.ts
+++ b/packages/server/src/integrations/tests/rest.spec.ts
@@ -29,6 +29,7 @@ import { default as RestIntegration } from "../rest"
 import { RestAuthType } from "@budibase/types"
 import fetch from "node-fetch"
 import { objectStoreTestProviders } from "./utils"
+import { Readable } from "stream"
 
 const FormData = require("form-data")
 const { URLSearchParams } = require("url")
@@ -638,7 +639,9 @@ describe("REST Integration", () => {
       const responseData = Buffer.from("teest file contnt")
       const filename = "test.tar.gz"
       const contentType = "application/gzip"
-
+      const mockReadable = new Readable()
+      mockReadable.push(responseData)
+      mockReadable.push(null)
       ;(fetch as unknown as jest.Mock).mockImplementationOnce(() =>
         Promise.resolve({
           headers: {
@@ -652,7 +655,7 @@ describe("REST Integration", () => {
                 return `attachment; filename="${filename}"`
             },
           },
-          arrayBuffer: jest.fn(() => Promise.resolve(responseData)),
+          body: mockReadable,
         })
       )
 
@@ -675,9 +678,10 @@ describe("REST Integration", () => {
 
     it("uploads file with non ascii filename to object store and returns signed URL ", async () => {
       const responseData = Buffer.from("teest file contnt")
-      const non_ascii_filename = "ex%C3%A4mple.txt"
       const contentType = "text/plain"
-
+      const mockReadable = new Readable()
+      mockReadable.push(responseData)
+      mockReadable.push(null)
       ;(fetch as unknown as jest.Mock).mockImplementationOnce(() =>
         Promise.resolve({
           headers: {
@@ -693,7 +697,7 @@ describe("REST Integration", () => {
                 return `attachment; filename="£ and ? rates.pdf"; filename*=UTF-8\'\'%C2%A3%20and%20%E2%82%AC%20rates.pdf`
             },
           },
-          arrayBuffer: jest.fn(() => Promise.resolve(responseData)),
+          body: mockReadable,
         })
       )
 
diff --git a/packages/server/src/integrations/utils.ts b/packages/server/src/integrations/utils.ts
index 9affc9917a..28ecda899c 100644
--- a/packages/server/src/integrations/utils.ts
+++ b/packages/server/src/integrations/utils.ts
@@ -468,7 +468,8 @@ export async function handleFileResponse(
   filename: string,
   startTime: number
 ) {
-  let presignedUrl
+  let presignedUrl,
+    size = 0
   const fileExtension = filename.includes(".")
     ? filename.split(".").slice(1).join(".")
     : ""
@@ -478,7 +479,19 @@ export async function handleFileResponse(
   const bucket = objectStore.ObjectStoreBuckets.TEMP
 
   const stream = response.body.pipe(bl((error, data) => data))
+
   if (response.body) {
+    const contentLength = response.headers.get("content-length")
+    if (contentLength) {
+      size = parseInt(contentLength, 10)
+    } else {
+      const chunks: Buffer[] = []
+      for await (const chunk of response.body) {
+        chunks.push(chunk)
+        size += chunk.length
+      }
+    }
+
     await objectStore.streamUpload({
       bucket,
       filename: key,
@@ -490,7 +503,7 @@ export async function handleFileResponse(
   presignedUrl = await objectStore.getPresignedUrl(bucket, key, 600)
   return {
     data: {
-      size: stream.byteLength,
+      size,
       name: processedFileName,
       url: presignedUrl,
       extension: fileExtension,
@@ -498,7 +511,7 @@ export async function handleFileResponse(
     },
     info: {
       code: response.status,
-      size: formatBytes(stream.byteLength),
+      size: formatBytes(size.toString()),
       time: `${Math.round(performance.now() - startTime)}ms`,
     },
   }