From 3cf4f2e4364dbf9675fa3c3683e1235bdbf23634 Mon Sep 17 00:00:00 2001
From: Chase Coalwell <782571+srchase@users.noreply.github.com>
Date: Wed, 21 Jun 2023 17:33:53 -0600
Subject: [PATCH] Add collect-stream-body

---
 packages/smithy-client/package.json           |  1 +
 .../src/collect-stream-body.spec.ts           | 49 +++++++++++++++++++
 .../smithy-client/src/collect-stream-body.ts  | 26 ++++++++++
 packages/smithy-client/src/index.ts           |  1 +
 yarn.lock                                     |  3 +-
 5 files changed, 79 insertions(+), 1 deletion(-)
 create mode 100644 packages/smithy-client/src/collect-stream-body.spec.ts
 create mode 100644 packages/smithy-client/src/collect-stream-body.ts

diff --git a/packages/smithy-client/package.json b/packages/smithy-client/package.json
index 7ab386fd1b5..f9e0a869136 100644
--- a/packages/smithy-client/package.json
+++ b/packages/smithy-client/package.json
@@ -25,6 +25,7 @@
   "dependencies": {
     "@smithy/middleware-stack": "workspace:^",
     "@smithy/types": "workspace:^",
+    "@smithy/util-stream": "workspace:^",
     "tslib": "^2.5.0"
   },
   "engines": {
diff --git a/packages/smithy-client/src/collect-stream-body.spec.ts b/packages/smithy-client/src/collect-stream-body.spec.ts
new file mode 100644
index 00000000000..58f911645d5
--- /dev/null
+++ b/packages/smithy-client/src/collect-stream-body.spec.ts
@@ -0,0 +1,49 @@
+import { Uint8ArrayBlobAdapter } from "@smithy/util-stream";
+
+import { collectBody } from "./collect-stream-body";
+
+describe(collectBody.name, () => {
+  it("passes through Uint8Array", async () => {
+    const body = new Uint8Array();
+    const arr = await collectBody(body, {
+      async streamCollector(stream: any) {
+        return new Uint8Array(stream);
+      },
+    });
+
+    expect(arr).toBe(body);
+  });
+
+  it("uses the contextual streamCollector", async () => {
+    const body = "x";
+    const arr = await collectBody(body, {
+      async streamCollector(stream: any) {
+        return Uint8ArrayBlobAdapter.fromString(stream);
+      },
+    });
+
+    expect(arr.transformToString()).toEqual("x");
+  });
+
+  it("uses the contextual streamCollector for empty string", async () => {
+    const body = "";
+    const arr = await collectBody(body, {
+      async streamCollector(stream: any) {
+        return Uint8ArrayBlobAdapter.fromString(stream);
+      },
+    });
+
+    expect(arr.transformToString()).toEqual("");
+  });
+
+  it("defaults to an empty Uint8Array", async () => {
+    const body = null;
+    const arr = await collectBody(body, {
+      async streamCollector(stream: any) {
+        return Uint8ArrayBlobAdapter.fromString(stream);
+      },
+    });
+
+    expect(arr.transformToString()).toEqual("");
+  });
+});
diff --git a/packages/smithy-client/src/collect-stream-body.ts b/packages/smithy-client/src/collect-stream-body.ts
new file mode 100644
index 00000000000..a20b5211b5e
--- /dev/null
+++ b/packages/smithy-client/src/collect-stream-body.ts
@@ -0,0 +1,26 @@
+import { SerdeContext } from "@smithy/types";
+import { Uint8ArrayBlobAdapter } from "@smithy/util-stream";
+
+/**
+ * @internal
+ *
+ * Collect low-level response body stream to Uint8Array.
+ */
+export const collectBody = async (
+  streamBody: any = new Uint8Array(),
+  context: {
+    streamCollector: SerdeContext["streamCollector"];
+  }
+): Promise<Uint8ArrayBlobAdapter> => {
+  if (streamBody instanceof Uint8Array) {
+    return Uint8ArrayBlobAdapter.mutate(streamBody);
+  }
+
+  if (!streamBody) {
+    return Uint8ArrayBlobAdapter.mutate(new Uint8Array());
+  }
+
+  const fromContext = context.streamCollector(streamBody);
+
+  return Uint8ArrayBlobAdapter.mutate(await fromContext);
+};
diff --git a/packages/smithy-client/src/index.ts b/packages/smithy-client/src/index.ts
index dc08b5e47e8..a8122e96193 100644
--- a/packages/smithy-client/src/index.ts
+++ b/packages/smithy-client/src/index.ts
@@ -1,5 +1,6 @@
 export * from "./NoOpLogger";
 export * from "./client";
+export * from "./collect-stream-body";
 export * from "./command";
 export * from "./constants";
 export * from "./create-aggregated-client";
diff --git a/yarn.lock b/yarn.lock
index 0a2d3a72c29..5123085eac9 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -2263,6 +2263,7 @@ __metadata:
   dependencies:
     "@smithy/middleware-stack": "workspace:^"
     "@smithy/types": "workspace:^"
+    "@smithy/util-stream": "workspace:^"
     "@tsconfig/recommended": 1.0.1
     "@types/node": ^14.14.31
     concurrently: 7.0.0
@@ -2517,7 +2518,7 @@ __metadata:
   languageName: unknown
   linkType: soft
 
-"@smithy/util-stream@workspace:packages/util-stream":
+"@smithy/util-stream@workspace:^, @smithy/util-stream@workspace:packages/util-stream":
   version: 0.0.0-use.local
   resolution: "@smithy/util-stream@workspace:packages/util-stream"
   dependencies: