fix(tmdt): stream 3D assets from S3 on init

awslabs · Jun 12, 2024 · ee0e11b · ee0e11b
1 parent 0f96c55
commit ee0e11b
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 29 deletions.
diff --git a/packages/tools-iottwinmaker/src/commands/init.spec.ts b/packages/tools-iottwinmaker/src/commands/init.spec.ts
@@ -52,6 +52,7 @@ beforeEach(() => {
   s3Mock.reset();
   twinmakerMock.reset();
   jest.spyOn(fs, 'writeFileSync');
+  jest.spyOn(fs, 'createWriteStream');
 });
 
 it('throws error when given tmdt project that does not exist', async () => {
@@ -148,33 +149,28 @@ it('creates a tmdt project with one model and one scene when given one scene and
     .resolves(emptyListComponentTypesResp);
   twinmakerMock.on(ListEntitiesCommand).resolves(emptyListEntitiesResp);
   twinmakerMock.on(ListScenesCommand).resolves(oneSceneListScenesResp);
+  const filePipeMock = jest.fn();
+  const s3ResponseBodyMock = (content: string) => {
+    return {
+      transformToString: () => {
+        return Promise.resolve(content);
+      },
+      pipe: filePipeMock,
+    } as unknown as SdkStream<Blob>;
+  };
+
   s3Mock
     .on(GetObjectCommand, { Bucket: 'workspace-bucket', Key: 'scene1.json' })
     .resolves({
       $metadata: {},
-      Body: {
-        transformToString: () => {
-          return Promise.resolve(JSON.stringify(scene1, null, 4));
-        },
-      } as SdkStream<Blob>,
+      Body: s3ResponseBodyMock(JSON.stringify(scene1, null, 4)),
     });
   s3Mock
     .on(GetObjectCommand, { Bucket: 'workspace-bucket', Key: 'model1.glb' })
     .resolves({
       $metadata: {},
-      Body: {
-        transformToString: () => {
-          return Promise.resolve(fakeModelData);
-        },
-      } as SdkStream<Blob>,
+      Body: s3ResponseBodyMock(fakeModelData),
     });
-  expect(
-    streamToBuffer({
-      transformToString: () => {
-        return Promise.resolve(JSON.stringify(scene1, null, 4));
-      },
-    } as SdkStream<Blob>)
-  ).toEqual('mock');
 
   const argv2 = {
     _: ['init'],
@@ -204,6 +200,7 @@ it('creates a tmdt project with one model and one scene when given one scene and
     `${outDir}/scene1.json`,
     JSON.stringify(scene1, null, 4)
   );
+  expect(filePipeMock).toBeCalledTimes(1);
 });
 
 it('creates a tmdt project with one entity when given one entity', async () => {
@@ -249,25 +246,27 @@ it('creates a fully populated tmdt project when given a full workspace', async (
   twinmakerMock.on(ListEntitiesCommand).resolves(oneEntityListEntitiesResp);
   twinmakerMock.on(GetEntityCommand).resolves(getEntity1Resp);
   twinmakerMock.on(ListScenesCommand).resolves(oneSceneListScenesResp);
+  const filePipeMock = jest.fn();
+  const s3ResponseBodyMock = (content: string) => {
+    return {
+      transformToString: () => {
+        return Promise.resolve(content);
+      },
+      pipe: filePipeMock,
+    } as unknown as SdkStream<Blob>;
+  };
+
   s3Mock
     .on(GetObjectCommand, { Bucket: 'workspace-bucket', Key: 'scene1.json' })
     .resolves({
       $metadata: {},
-      Body: {
-        transformToString: () => {
-          return Promise.resolve(JSON.stringify(scene1, null, 4));
-        },
-      } as SdkStream<Blob>,
+      Body: s3ResponseBodyMock(JSON.stringify(scene1, null, 4)),
     });
   s3Mock
     .on(GetObjectCommand, { Bucket: 'workspace-bucket', Key: 'model1.glb' })
     .resolves({
       $metadata: {},
-      Body: {
-        transformToString: () => {
-          return Promise.resolve(fakeModelData);
-        },
-      } as SdkStream<Blob>,
+      Body: s3ResponseBodyMock(fakeModelData),
     });
   expect(
     streamToBuffer({
@@ -309,4 +308,5 @@ it('creates a fully populated tmdt project when given a full workspace', async (
     `${outDir}/scene1.json`,
     JSON.stringify(scene1, null, 4)
   );
+  expect(filePipeMock).toBeCalledTimes(1);
 });
diff --git a/packages/tools-iottwinmaker/src/commands/init.ts b/packages/tools-iottwinmaker/src/commands/init.ts
@@ -16,6 +16,7 @@ import {
 } from '@aws-sdk/client-iottwinmaker';
 import { verifyWorkspaceExists } from '../lib/utils';
 import { streamToBuffer } from '../lib/stream-to-buffer';
+import { Readable } from 'stream';
 
 export type Options = {
   region: string;
@@ -256,13 +257,19 @@ async function import_scenes_and_models(
           }
         }
 
+        // Stream S3 content and write to disk
         const data = await aws().s3.getObject({ Bucket: s3bucket, Key: s3key });
-        const bodyContents = (await streamToBuffer(data.Body)) as Buffer;
-        fs.writeFileSync(path.join(outDir, '3d_models', s3key), bodyContents);
+        const writeStream = fs.createWriteStream(
+          path.join(outDir, '3d_models', s3key)
+        );
+        (data.Body as Readable).pipe(writeStream);
         tmdt_config['models'].push(s3key);
 
         // handle binary data references in gltf files - https://www.khronos.org/files/gltf20-reference-guide.pdf
         if (s3key.endsWith('.gltf')) {
+          // Load S3 content in memory
+          const bodyContents = (await streamToBuffer(data.Body)) as Buffer;
+
           const gltfData = JSON.parse(bodyContents.toString('utf-8'));
           if (gltfData['buffers']) {
             for (const buffer of gltfData['buffers']) {