ec2: improve QuickPick UI (#3596)

* add empty connect command

* create sample QuickPick on click

* hacky prompt + func to list instanceId

* very hacky way to create prompt with instanceIds

* move extractInstanceIds to own function

* utilize lastTouchedRegion in prompt

* throw cancellationError on user cancel

* add a title to quickpick

* refactor: increase testability

* move test utility func to shared utility file

* add general test cases for extractInstanceIdsFromReservations

* add basic test for prompter

* configure command to devMode

* refactor to utilize wizard + integrate regionSubmenu

* add new testing file

* refactor to utilize regionSubmenu in isolation, rather than wrapped in Wizard class

* fix awkward indent

* add test file from feature/cwl branch

* remove old prompter old + tests

* delete tests that rely on feature/cwl changes

* refactor to avoid circular dependency

* fix improper headers + imports

* introduce Ec2ConnectClient to handle connection

* remove dead parameter

* close connection on terminal close

* add a little more to error msg

* remove dead imports

* log error to console

* rename function

Co-authored-by: Justin M. Keyes <[email protected]>

* fix header

Co-authored-by: Justin M. Keyes <[email protected]>

* fix headers and imports

* remove year from header

* delete outdated test case

* delete outdated test file

* remove year from copyright header

* fix formatting of files

* remove alias in Ec2Instance

* utilize interface for object-like shape

* improve style + start of work on debugging error

* make distinction between status error and permission error

* start to add tests for error handling

* move code to general Ec2Client

* change Log Groups to selection in region submenu

* refactor to avoid circular dependency

* fix formatting

* fix formatting

* comment out the ec2Client we are not using

* style fix

* fix formatting

* add test for error handling

* fix formatting

* add extra wrapper

* generalize some code to a remoteSession file

* add space between pieces

* retrieve IAM role attached to instance when fails to connect

* remove dead import

* refactor to avoid circular dependency

* checks if relevant policies are on attached role

* use both probes to determine source of error

* update tests

* change callback to async/await

* remove onError parameter and utilize .catch instead

* add test for permissions detection

* fix test

* make error handling a try-catch

* swap ec2Client to v3

* switch over SSM client

* update tests and signatures for new error type in v3

* fix formatting

* remove dead import

* avoid unnecessary default prefix on client name

Co-authored-by: Justin M. Keyes <[email protected]>

* fix old import

* refactor defaultEc2client -> ec2Client

* rename Ec2ConnectClient

* fix outdated imports in the test

* add icons and detail to description

* update tests to use sdk3

* update tests and function names

* bubble error up

* update test case

* Remove please in text

Co-authored-by: Justin M. Keyes <[email protected]>

* Remove please in text in other error

Co-authored-by: Justin M. Keyes <[email protected]>

* remove unnecessary types and update tests to match

* add match on detail, and instance id label

* undo migration within package files and high-level model file

* update tests to downgrade

* undo some minor tweaks to the clients

* use v2 import in the test file

* display no-name instead of id when name not present

* change Ec2Instance to extends Ec2.Instance

* use 'get' over 'extract' in language.

Co-authored-by: Justin M. Keyes <[email protected]>

* update references to getInstancesFromReservations

---------

Co-authored-by: Justin M. Keyes <[email protected]>

Author: Hweinstock

src/ec2/prompter.ts

@@ -4,13 +4,15 @@
  */
 
 import { RegionSubmenu, RegionSubmenuResponse } from '../shared/ui/common/regionSubmenu'
-import { Ec2Selection, getInstanceIdsFromRegion } from './utils'
+import { Ec2Selection, getInstancesFromRegion } from './utils'
 import { DataQuickPickItem } from '../shared/ui/pickerPrompter'
+import { Ec2Instance } from '../shared/clients/ec2Client'
 
-function asQuickpickItem(instanceId: string): DataQuickPickItem<string> {
+function asQuickpickItem(instance: Ec2Instance): DataQuickPickItem<string> {
     return {
-        label: instanceId,
-        data: instanceId,
+        label: '$(terminal) \t' + (instance.name ?? '(no name)'),
+        detail: instance.InstanceId,
+        data: instance.InstanceId,
     }
 }
 
@@ -23,8 +25,8 @@ export function handleEc2ConnectPrompterResponse(response: RegionSubmenuResponse
 
 export function createEc2ConnectPrompter(): RegionSubmenu<string> {
     return new RegionSubmenu(
-        async region => (await getInstanceIdsFromRegion(region)).map(asQuickpickItem).promise(),
-        { title: 'Select EC2 Instance Id' },
+        async region => (await getInstancesFromRegion(region)).map(asQuickpickItem).promise(),
+        { title: 'Select EC2 Instance Id', matchOnDetail: true },
         { title: 'Select Region for EC2 Instance' },
         'Instances'
     )

src/ec2/utils.ts

@@ -4,14 +4,14 @@
  */
 
 import { AsyncCollection } from '../shared/utilities/asyncCollection'
-import { Ec2Client } from '../shared/clients/ec2Client'
+import { Ec2Client, Ec2Instance } from '../shared/clients/ec2Client'
 
 export interface Ec2Selection {
     instanceId: string
     region: string
 }
 
-export async function getInstanceIdsFromRegion(regionCode: string): Promise<AsyncCollection<string>> {
+export async function getInstancesFromRegion(regionCode: string): Promise<AsyncCollection<Ec2Instance>> {
     const client = new Ec2Client(regionCode)
-    return await client.getInstanceIds()
+    return await client.getInstances()
 }

src/shared/clients/ec2Client.ts

@@ -4,36 +4,44 @@
  */
 
 import { EC2 } from 'aws-sdk'
-import globals from '../extensionGlobals'
 import { AsyncCollection } from '../utilities/asyncCollection'
 import { pageableToCollection } from '../utilities/collectionUtils'
 import { IamInstanceProfile } from 'aws-sdk/clients/ec2'
+import globals from '../extensionGlobals'
+
+export interface Ec2Instance extends EC2.Instance {
+    name?: string
+}
 
 export class Ec2Client {
     public constructor(public readonly regionCode: string) {}
 
     private async createSdkClient(): Promise<EC2> {
         return await globals.sdkClientBuilder.createAwsService(EC2, undefined, this.regionCode)
     }
-    public async getInstanceIds(): Promise<AsyncCollection<string>> {
+    public async getInstances(): Promise<AsyncCollection<Ec2Instance>> {
         const client = await this.createSdkClient()
         const requester = async (request: EC2.DescribeInstancesRequest) => client.describeInstances(request).promise()
+        const collection = pageableToCollection(requester, {}, 'NextToken', 'Reservations')
+        const instances = this.getInstancesFromReservations(collection)
+        return instances
+    }
 
-        const instanceIds = this.extractInstanceIdsFromReservations(
-            pageableToCollection(requester, {}, 'NextToken', 'Reservations')
-        )
-        return instanceIds
+    private lookupTagKey(tags: EC2.Tag[], targetKey: string): string | undefined {
+        return tags.filter(tag => tag.Key == targetKey)[0].Value
     }
 
-    public extractInstanceIdsFromReservations(
+    public getInstancesFromReservations(
         reservations: AsyncCollection<EC2.ReservationList | undefined>
-    ): AsyncCollection<string> {
+    ): AsyncCollection<Ec2Instance> {
         return reservations
             .flatten()
             .map(instanceList => instanceList?.Instances)
             .flatten()
-            .map(instance => instance?.InstanceId)
-            .filter(instanceId => instanceId !== undefined)
+            .filter(instance => instance!.InstanceId !== undefined)
+            .map(instance => {
+                return instance!.Tags ? { ...instance, name: this.lookupTagKey(instance!.Tags!, 'Name') } : instance!
+            })
     }
 
     public async getInstanceStatus(instanceId: string): Promise<EC2.InstanceStateName> {

src/shared/clients/ssmClient.ts

@@ -4,9 +4,9 @@
  */
 
 import { AWSError, SSM } from 'aws-sdk'
-import globals from '../extensionGlobals'
-import { PromiseResult } from 'aws-sdk/lib/request'
 import { getLogger } from '../logger/logger'
+import { PromiseResult } from 'aws-sdk/lib/request'
+import globals from '../extensionGlobals'
 
 export class SsmClient {
     public constructor(public readonly regionCode: string) {}
@@ -26,7 +26,8 @@ export class SsmClient {
             .catch(err => {
                 getLogger().warn(`ssm: failed to terminate session "${sessionId}": %s`, err)
             })
-        return termination
+
+        return termination!
     }
 
     public async startSession(target: string): Promise<PromiseResult<SSM.StartSessionResponse, AWSError>> {

src/test/ec2/model.test.ts

@@ -7,10 +7,10 @@ import * as assert from 'assert'
 import { Ec2ConnectErrorCode, Ec2ConnectionManager } from '../../ec2/model'
 import { SsmClient } from '../../shared/clients/ssmClient'
 import { Ec2Client } from '../../shared/clients/ec2Client'
-import { AWSError } from 'aws-sdk'
 import { attachedPoliciesListType } from 'aws-sdk/clients/iam'
 import { Ec2Selection } from '../../ec2/utils'
 import { ToolkitError } from '../../shared/errors'
+import { AWSError, EC2 } from 'aws-sdk'
 
 describe('Ec2ConnectClient', function () {
     class MockSsmClient extends SsmClient {
@@ -24,8 +24,8 @@ describe('Ec2ConnectClient', function () {
             super('test-region')
         }
 
-        public override async getInstanceStatus(instanceId: string): Promise<string> {
-            return instanceId.split(':')[0]
+        public override async getInstanceStatus(instanceId: string): Promise<EC2.InstanceStateName> {
+            return instanceId.split(':')[0] as EC2.InstanceStateName
         }
     }
 
@@ -44,7 +44,7 @@ describe('Ec2ConnectClient', function () {
     }
     describe('handleStartSessionError', async function () {
         let client: MockEc2ConnectClientForError
-        let dummyError: AWSError
+        const dummyError: AWSError = { name: 'testName', message: 'testMessage', code: 'testCode', time: new Date() }
 
         class MockEc2ConnectClientForError extends MockEc2ConnectClient {
             public override async hasProperPolicies(instanceId: string): Promise<boolean> {
@@ -53,18 +53,17 @@ describe('Ec2ConnectClient', function () {
         }
         before(function () {
             client = new MockEc2ConnectClientForError()
-            dummyError = {} as AWSError
         })
 
         it('determines which error to throw based on if instance is running', async function () {
             async function assertThrowsErrorCode(testInstance: Ec2Selection, errCode: Ec2ConnectErrorCode) {
                 try {
                     await client.handleStartSessionError(dummyError, testInstance)
-                    throw new assert.AssertionError({ message: "Didn't throw any error" })
                 } catch (err: unknown) {
                     assert.strictEqual((err as ToolkitError).code, errCode)
                 }
             }
+
             await assertThrowsErrorCode(
                 {
                     instanceId: 'pending:noPolicies',

src/test/shared/clients/defaultEc2Client.test.ts

@@ -5,16 +5,16 @@
 
 import * as assert from 'assert'
 import { AsyncCollection } from '../../../shared/utilities/asyncCollection'
-import { EC2 } from 'aws-sdk'
 import { toCollection } from '../../../shared/utilities/asyncCollection'
 import { intoCollection } from '../../../shared/utilities/collectionUtils'
 import { Ec2Client } from '../../../shared/clients/ec2Client'
+import { EC2 } from 'aws-sdk'
 
-describe('extractInstanceIdsFromReservations', function () {
+describe('extractInstancesFromReservations', function () {
     const client = new Ec2Client('')
     it('returns empty when given empty collection', async function () {
         const actualResult = await client
-            .extractInstanceIdsFromReservations(
+            .getInstancesFromReservations(
                 toCollection(async function* () {
                     yield []
                 }) as AsyncCollection<EC2.ReservationList>
@@ -30,27 +30,37 @@ describe('extractInstanceIdsFromReservations', function () {
                 Instances: [
                     {
                         InstanceId: 'id1',
+                        Tags: [{ Key: 'Name', Value: 'name1' }],
                     },
                     {
                         InstanceId: 'id2',
+                        Tags: [{ Key: 'Name', Value: 'name2' }],
                     },
                 ],
             },
             {
                 Instances: [
                     {
                         InstanceId: 'id3',
+                        Tags: [{ Key: 'Name', Value: 'name3' }],
                     },
                     {
                         InstanceId: 'id4',
+                        Tags: [{ Key: 'Name', Value: 'name4' }],
                     },
                 ],
             },
         ]
-        const actualResult = await client
-            .extractInstanceIdsFromReservations(intoCollection([testReservationsList]))
-            .promise()
-        assert.deepStrictEqual(['id1', 'id2', 'id3', 'id4'], actualResult)
+        const actualResult = await client.getInstancesFromReservations(intoCollection([testReservationsList])).promise()
+        assert.deepStrictEqual(
+            [
+                { InstanceId: 'id1', name: 'name1', Tags: [{ Key: 'Name', Value: 'name1' }] },
+                { InstanceId: 'id2', name: 'name2', Tags: [{ Key: 'Name', Value: 'name2' }] },
+                { InstanceId: 'id3', name: 'name3', Tags: [{ Key: 'Name', Value: 'name3' }] },
+                { InstanceId: 'id4', name: 'name4', Tags: [{ Key: 'Name', Value: 'name4' }] },
+            ],
+            actualResult
+        )
     }),
         // Unsure if this test case is needed, but the return type in the SDK makes it possible these are unknown/not returned.
         it('handles undefined and missing pieces in the ReservationList.', async function () {
@@ -69,14 +79,18 @@ describe('extractInstanceIdsFromReservations', function () {
                     Instances: [
                         {
                             InstanceId: 'id3',
+                            Tags: [{ Key: 'Name', Value: 'name3' }],
                         },
                         {},
                     ],
                 },
             ]
             const actualResult = await client
-                .extractInstanceIdsFromReservations(intoCollection([testReservationsList]))
+                .getInstancesFromReservations(intoCollection([testReservationsList]))
                 .promise()
-            assert.deepStrictEqual(['id1', 'id3'], actualResult)
+            assert.deepStrictEqual(
+                [{ InstanceId: 'id1' }, { InstanceId: 'id3', name: 'name3', Tags: [{ Key: 'Name', Value: 'name3' }] }],
+                actualResult
+            )
         })
 })