fix(test): fix flaky test related to ec2.updateStatus

packages/core/src/awsService/ec2/explorer/ec2InstanceNode.ts

@@ -29,6 +29,7 @@ export class Ec2InstanceNode extends AWSTreeNodeBase implements AWSResourceNode
         public readonly instance: SafeEc2Instance
     ) {
         super('')
+        this.parent.addChild(this)
         this.updateInstance(instance)
         this.id = this.InstanceId
     }
@@ -41,7 +42,7 @@ export class Ec2InstanceNode extends AWSTreeNodeBase implements AWSResourceNode
         this.tooltip = `${this.name}\n${this.InstanceId}\n${this.instance.LastSeenStatus}\n${this.arn}`
 
         if (this.isPending()) {
-            this.parent.pollingSet.start(this.InstanceId)
+            this.parent.trackPendingNode(this.InstanceId)
         }
     }
 

packages/core/src/awsService/ec2/explorer/ec2ParentNode.ts

@@ -41,6 +41,13 @@ export class Ec2ParentNode extends AWSTreeNodeBase {
         })
     }
 
+    public trackPendingNode(instanceId: string) {
+        if (!this.ec2InstanceNodes.has(instanceId)) {
+            throw new Error(`Attempt to track ec2 node ${instanceId} that isn't a child`)
+        }
+        this.pollingSet.start(instanceId)
+    }
+
     public async updateChildren(): Promise<void> {
         const ec2Instances = await (await this.ec2Client.getInstances()).toMap((instance) => instance.InstanceId)
         updateInPlace(
@@ -52,9 +59,18 @@ export class Ec2ParentNode extends AWSTreeNodeBase {
         )
     }
 
+    public getInstanceNode(instanceId: string): Ec2InstanceNode {
+        const childNode = this.ec2InstanceNodes.get(instanceId)
+        if (childNode) {
+            return childNode
+        } else {
+            throw new Error(`Node with id ${instanceId} from polling set not found`)
+        }
+    }
+
     private async updatePendingNodes() {
         for (const instanceId of this.pollingSet.values()) {
-            const childNode = this.ec2InstanceNodes.get(instanceId)!
+            const childNode = this.getInstanceNode(instanceId)
             await this.updatePendingNode(childNode)
         }
     }
@@ -71,6 +87,10 @@ export class Ec2ParentNode extends AWSTreeNodeBase {
         this.ec2InstanceNodes = new Map<string, Ec2InstanceNode>()
     }
 
+    public addChild(node: Ec2InstanceNode) {
+        this.ec2InstanceNodes.set(node.InstanceId, node)
+    }
+
     public async refreshNode(): Promise<void> {
         await this.clearChildren()
         await vscode.commands.executeCommand('aws.refreshAwsExplorerNode', this)

packages/core/src/test/awsService/ec2/explorer/ec2InstanceNode.test.ts

@@ -12,6 +12,8 @@ import {
 } from '../../../../awsService/ec2/explorer/ec2InstanceNode'
 import { Ec2Client, SafeEc2Instance, getNameOfInstance } from '../../../../shared/clients/ec2Client'
 import { Ec2ParentNode } from '../../../../awsService/ec2/explorer/ec2ParentNode'
+import * as sinon from 'sinon'
+import { PollingSet } from '../../../../shared/utilities/pollingSet'
 
 describe('ec2InstanceNode', function () {
     let testNode: Ec2InstanceNode
@@ -30,12 +32,16 @@ describe('ec2InstanceNode', function () {
             ],
             LastSeenStatus: 'running',
         }
+        sinon.stub(Ec2InstanceNode.prototype, 'updateStatus')
+        // Don't want to be polling here, that is tested in ../ec2ParentNode.test.ts
+        // disabled here for convenience (avoiding race conditions with timeout)
+        sinon.stub(PollingSet.prototype, 'start')
         const testClient = new Ec2Client('')
         const testParentNode = new Ec2ParentNode(testRegion, testPartition, testClient)
         testNode = new Ec2InstanceNode(testParentNode, testClient, 'testRegion', 'testPartition', testInstance)
     })
 
-    this.beforeEach(function () {
+    beforeEach(function () {
         testNode.updateInstance(testInstance)
     })
 

packages/core/src/test/awsService/ec2/explorer/ec2ParentNode.test.ts

@@ -17,17 +17,14 @@ import { EC2 } from 'aws-sdk'
 import { AsyncCollection } from '../../../../shared/utilities/asyncCollection'
 import * as FakeTimers from '@sinonjs/fake-timers'
 import { installFakeClock } from '../../../testUtil'
-import { PollingSet } from '../../../../shared/utilities/pollingSet'
 
 describe('ec2ParentNode', function () {
     let testNode: Ec2ParentNode
-    let defaultInstances: SafeEc2Instance[]
     let client: Ec2Client
     let getInstanceStub: sinon.SinonStub<[filters?: EC2.Filter[] | undefined], Promise<AsyncCollection<EC2.Instance>>>
     let clock: FakeTimers.InstalledClock
     let refreshStub: sinon.SinonStub<[], Promise<void>>
-    let clearTimerStub: sinon.SinonStub<[], void>
-
+    let statusUpdateStub: sinon.SinonStub<[status: string], Promise<string>>
     const testRegion = 'testRegion'
     const testPartition = 'testPartition'
 
@@ -45,36 +42,19 @@ describe('ec2ParentNode', function () {
         client = new Ec2Client(testRegion)
         clock = installFakeClock()
         refreshStub = sinon.stub(Ec2InstanceNode.prototype, 'refreshNode')
-        clearTimerStub = sinon.stub(PollingSet.prototype, 'clearTimer')
-        defaultInstances = [
-            { Name: 'firstOne', InstanceId: '0', LastSeenStatus: 'running' },
-            { Name: 'secondOne', InstanceId: '1', LastSeenStatus: 'running' },
-        ]
+        statusUpdateStub = sinon.stub(Ec2Client.prototype, 'getInstanceStatus')
     })
 
     beforeEach(function () {
         getInstanceStub = sinon.stub(Ec2Client.prototype, 'getInstances')
-        defaultInstances = [
-            { Name: 'firstOne', InstanceId: '0', LastSeenStatus: 'running' },
-            { Name: 'secondOne', InstanceId: '1', LastSeenStatus: 'stopped' },
-        ]
-
-        getInstanceStub.callsFake(async () =>
-            intoCollection(
-                defaultInstances.map((instance) => ({
-                    InstanceId: instance.InstanceId,
-                    Tags: [{ Key: 'Name', Value: instance.Name }],
-                }))
-            )
-        )
-
         testNode = new Ec2ParentNode(testRegion, testPartition, client)
         refreshStub.resetHistory()
-        clearTimerStub.resetHistory()
     })
 
     afterEach(function () {
         getInstanceStub.restore()
+        testNode.pollingSet.clear()
+        testNode.pollingSet.clearTimer()
     })
 
     after(function () {
@@ -91,10 +71,14 @@ describe('ec2ParentNode', function () {
     })
 
     it('has instance child nodes', async function () {
-        getInstanceStub.resolves(mapToInstanceCollection(defaultInstances))
+        const instances = [
+            { Name: 'firstOne', InstanceId: '0', LastSeenStatus: 'running' },
+            { Name: 'secondOne', InstanceId: '1', LastSeenStatus: 'stopped' },
+        ]
+        getInstanceStub.resolves(mapToInstanceCollection(instances))
         const childNodes = await testNode.getChildren()
 
-        assert.strictEqual(childNodes.length, defaultInstances.length, 'Unexpected child count')
+        assert.strictEqual(childNodes.length, instances.length, 'Unexpected child count')
 
         childNodes.forEach((node) =>
             assert.ok(node instanceof Ec2InstanceNode, 'Expected child node to be Ec2InstanceNode')
@@ -151,14 +135,13 @@ describe('ec2ParentNode', function () {
         ]
 
         getInstanceStub.resolves(mapToInstanceCollection(instances))
-
         await testNode.updateChildren()
         assert.strictEqual(testNode.pollingSet.size, 1)
         getInstanceStub.restore()
     })
 
     it('does not refresh explorer when timer goes off if status unchanged', async function () {
-        const statusUpdateStub = sinon.stub(Ec2Client.prototype, 'getInstanceStatus').resolves('pending')
+        statusUpdateStub = statusUpdateStub.resolves('pending')
         const instances = [
             { Name: 'firstOne', InstanceId: '0', LastSeenStatus: 'pending' },
             { Name: 'secondOne', InstanceId: '1', LastSeenStatus: 'stopped' },
@@ -170,16 +153,56 @@ describe('ec2ParentNode', function () {
         await testNode.updateChildren()
         await clock.tickAsync(6000)
         sinon.assert.notCalled(refreshStub)
-        statusUpdateStub.restore()
         getInstanceStub.restore()
     })
 
     it('does refresh explorer when timer goes and status changed', async function () {
+        statusUpdateStub = statusUpdateStub.resolves('running')
+        const instances = [{ Name: 'firstOne', InstanceId: '0', LastSeenStatus: 'pending' }]
+
+        getInstanceStub.resolves(mapToInstanceCollection(instances))
+        await testNode.updateChildren()
+
         sinon.assert.notCalled(refreshStub)
-        const statusUpdateStub = sinon.stub(Ec2Client.prototype, 'getInstanceStatus').resolves('running')
-        testNode.pollingSet.add('0')
         await clock.tickAsync(6000)
         sinon.assert.called(refreshStub)
-        statusUpdateStub.restore()
+    })
+
+    it('returns the node when in the map', async function () {
+        const instances = [{ Name: 'firstOne', InstanceId: 'node1', LastSeenStatus: 'pending' }]
+
+        getInstanceStub.resolves(mapToInstanceCollection(instances))
+        await testNode.updateChildren()
+        const node = testNode.getInstanceNode('node1')
+        assert.strictEqual(node.InstanceId, instances[0].InstanceId)
+        getInstanceStub.restore()
+    })
+
+    it('throws error when node not in map', async function () {
+        const instances = [{ Name: 'firstOne', InstanceId: 'node1', LastSeenStatus: 'pending' }]
+
+        getInstanceStub.resolves(mapToInstanceCollection(instances))
+        await testNode.updateChildren()
+        assert.throws(() => testNode.getInstanceNode('node2'))
+        getInstanceStub.restore()
+    })
+
+    it('adds node to polling set when asked to track it', async function () {
+        const instances = [{ Name: 'firstOne', InstanceId: 'node1', LastSeenStatus: 'pending' }]
+
+        getInstanceStub.resolves(mapToInstanceCollection(instances))
+        await testNode.updateChildren()
+        testNode.trackPendingNode('node1')
+        assert.strictEqual(testNode.pollingSet.size, 1)
+        getInstanceStub.restore()
+    })
+
+    it('throws error when asked to track non-child node', async function () {
+        const instances = [{ Name: 'firstOne', InstanceId: 'node1', LastSeenStatus: 'pending' }]
+
+        getInstanceStub.resolves(mapToInstanceCollection(instances))
+        await testNode.updateChildren()
+        assert.throws(() => testNode.trackPendingNode('node2'))
+        getInstanceStub.restore()
     })
 })