Merge a9ad134 into mainline

awslabs · Mar 3, 2021 · 92b1959 · 92b1959
2 parents b535b38 + a9ad134
commit 92b1959
Show file tree

Hide file tree

Showing 2 changed files with 104 additions and 3 deletions.
diff --git a/...se-raas-services/lib/environment/service-catalog/__tests__/environment-sc-service.test.js b/...se-raas-services/lib/environment/service-catalog/__tests__/environment-sc-service.test.js
@@ -644,6 +644,108 @@ describe('EnvironmentSCService', () => {
       ]);
       AWSMock.restore();
     });
+
+    it('Should call listNotebookInstances with NextToken if NextToken is returned', async () => {
+      // BUILD
+      indexesService.list = jest
+        .fn()
+        .mockResolvedValue([{ id: 'some-index-id', awsAccountId: 'some-aws-account-uuid' }]);
+      awsAccountsService.list = jest.fn().mockResolvedValue([
+        {
+          roleArn: 'some-role-arn',
+          externalId: 'some-external-id',
+          id: 'some-aws-account-uuid',
+          accountId: 'some-aws-account-id',
+        },
+      ]);
+      dbService.table.scan.mockResolvedValueOnce([
+        {
+          id: 'some-environment-id',
+          indexId: 'some-index-id',
+          status: 'STOPPING',
+          outputs: [{ OutputKey: 'NotebookInstanceName', OutputValue: 'notebook-instance-name' }],
+        },
+        {
+          id: 'some-environment-id-1',
+          indexId: 'some-index-id',
+          status: 'STOPPING',
+          outputs: [{ OutputKey: 'NotebookInstanceName', OutputValue: 'notebook-instance-name-1' }],
+        },
+        {
+          id: 'some-environment-id-2',
+          indexId: 'some-index-id',
+          status: 'STOPPING',
+          outputs: [{ OutputKey: 'NotebookInstanceName', OutputValue: 'notebook-instance-name-2' }],
+        },
+      ]);
+      service.update = jest.fn();
+
+      AWSMock.setSDKInstance(aws.sdk);
+      AWSMock.mock('STS', 'assumeRole', { Credentials: stsResponse });
+      const mockListNotebookInstance = jest
+        .fn()
+        .mockImplementationOnce((params, callback) => {
+          callback(null, {
+            NotebookInstances: [
+              {
+                NotebookInstanceName: 'notebook-instance-name',
+                NotebookInstanceStatus: 'Stopped',
+              },
+              {
+                NotebookInstanceName: 'notebook-instance-name-1',
+                NotebookInstanceStatus: 'InService',
+              },
+            ],
+            NextToken: 'some-next-token',
+          });
+        })
+        .mockImplementationOnce((params, callback) => {
+          if (params.NextToken === 'some-next-token') {
+            callback(null, {
+              NotebookInstances: [
+                {
+                  NotebookInstanceName: 'notebook-instance-name-2',
+                  NotebookInstanceStatus: 'Updating',
+                },
+              ],
+            });
+          } else {
+            callback({ message: 'NextToken is different from expected' }, null);
+          }
+        })
+        .mockImplementationOnce((params, callback) => {
+          callback({ message: `list notebook was called the third time, only 2 calls expected` }, null);
+        });
+      AWSMock.mock('SageMaker', 'listNotebookInstances', mockListNotebookInstance);
+
+      // OPERATE
+      const result = await service.pollAndSyncWsStatus(requestContext);
+      // CHECK
+      expect(result).toMatchObject([
+        {
+          accountId: 'some-aws-account-id',
+          ec2Updated: {},
+          sagemakerUpdated: {
+            'notebook-instance-name': {
+              currentStatus: 'STOPPED',
+              ddbID: 'some-environment-id',
+              staleStatus: 'STOPPING',
+            },
+            'notebook-instance-name-1': {
+              currentStatus: 'COMPLETED',
+              ddbID: 'some-environment-id-1',
+              staleStatus: 'STOPPING',
+            },
+            'notebook-instance-name-2': {
+              currentStatus: 'STARTING',
+              ddbID: 'some-environment-id-2',
+              staleStatus: 'STOPPING',
+            },
+          },
+        },
+      ]);
+      AWSMock.restore();
+    });
   });
 
   describe('changeWorkspaceRunState function', () => {

diff --git a/...aas/packages/base-raas-services/lib/environment/service-catalog/environment-sc-service.js b/...aas/packages/base-raas-services/lib/environment/service-catalog/environment-sc-service.js
@@ -234,11 +234,10 @@ class EnvironmentScService extends Service {
   async pollSageMakerRealtimeStatus(roleArn, externalId) {
     const aws = await this.service('aws');
     const sagemakerClient = await aws.getClientSdkForRole({ roleArn, externalId, clientName: 'SageMaker' });
-    const params = {};
+    const params = { MaxResults: 100 };
     const sagemakerRealtimeStatus = {};
-    let data;
     do {
-      data = await sagemakerClient.listNotebookInstances().promise(); // eslint-disable-line no-await-in-loop
+      const data = await sagemakerClient.listNotebookInstances(params).promise(); // eslint-disable-line no-await-in-loop
       params.NextToken = data.NextToken;
       data.NotebookInstances.forEach(instance => {
         sagemakerRealtimeStatus[instance.NotebookInstanceName] = instance.NotebookInstanceStatus;