Refactor sessionState tests

Author: siakmun-aws

packages/core/src/test/amazonq/session/testSetup.ts

@@ -5,6 +5,7 @@
 
 import sinon from 'sinon'
 import { createBasicTestConfig, createMockSessionStateConfig, TestSessionMocks } from '../utils'
+import { SessionStateConfig } from '../../../amazonq'
 
 export function createSessionTestSetup() {
     const conversationId = 'conversation-id'
@@ -20,7 +21,7 @@ export function createSessionTestSetup() {
     }
 }
 
-export async function beforeEachFunc(
+export async function createTestConfig(
     testMocks: TestSessionMocks,
     conversationId: string,
     uploadId: string,
@@ -33,3 +34,41 @@ export async function beforeEachFunc(
     const testConfig = createMockSessionStateConfig(basicConfig, testMocks)
     return testConfig
 }
+
+export interface TestContext {
+    conversationId: string
+    uploadId: string
+    tabId: string
+    currentCodeGenerationId: string
+    testConfig: SessionStateConfig
+    testMocks: Record<string, any>
+}
+
+export function createTestContext(): TestContext {
+    const { conversationId, uploadId, tabId, currentCodeGenerationId } = createSessionTestSetup()
+
+    return {
+        conversationId,
+        uploadId,
+        tabId,
+        currentCodeGenerationId,
+        testConfig: {} as SessionStateConfig,
+        testMocks: {},
+    }
+}
+
+export function setupTestHooks(context: TestContext) {
+    beforeEach(async () => {
+        context.testMocks = {}
+        context.testConfig = await createTestConfig(
+            context.testMocks,
+            context.conversationId,
+            context.uploadId,
+            context.currentCodeGenerationId
+        )
+    })
+
+    afterEach(() => {
+        sinon.restore()
+    })
+}

packages/core/src/test/amazonqDoc/session/sessionState.test.ts

@@ -6,35 +6,23 @@
 import * as vscode from 'vscode'
 import assert from 'assert'
 import sinon from 'sinon'
-import { DocPrepareCodeGenState, SessionStateConfig } from '../../../amazonqDoc'
+import { DocPrepareCodeGenState } from '../../../amazonqDoc'
 import { createMockSessionStateAction } from '../../amazonq/utils'
 
-import { TestSessionMocks } from '../../amazonq/utils'
-import { beforeEachFunc, createSessionTestSetup } from '../../amazonq/session/testSetup'
-
-let testMocks: TestSessionMocks
+import { createTestContext, setupTestHooks } from '../../amazonq/session/testSetup'
 
 describe('sessionStateDoc', () => {
-    const { conversationId, uploadId, tabId, currentCodeGenerationId } = createSessionTestSetup()
-    let testConfig: SessionStateConfig
-
-    beforeEach(async () => {
-        testMocks = {}
-        testConfig = await beforeEachFunc(testMocks, conversationId, uploadId, currentCodeGenerationId)
-    })
-
-    afterEach(() => {
-        sinon.restore()
-    })
+    const context = createTestContext()
+    setupTestHooks(context)
 
     describe('DocPrepareCodeGenState', () => {
         it('error when failing to prepare repo information', async () => {
             sinon.stub(vscode.workspace, 'findFiles').throws()
-            testMocks.createUploadUrl!.resolves({ uploadId: '', uploadUrl: '' })
+            context.testMocks.createUploadUrl!.resolves({ uploadId: '', uploadUrl: '' })
             const testAction = createMockSessionStateAction()
 
             await assert.rejects(() => {
-                return new DocPrepareCodeGenState(testConfig, [], [], [], tabId, 0).interact(testAction)
+                return new DocPrepareCodeGenState(context.testConfig, [], [], [], context.tabId, 0).interact(testAction)
             })
         })
     })

packages/core/src/test/amazonqFeatureDev/session/sessionState.test.ts

@@ -11,34 +11,21 @@ import {
     FeatureDevPrepareCodeGenState,
     FeatureDevCodeGenState,
 } from '../../../amazonqFeatureDev/session/sessionState'
-import { SessionStateConfig } from '../../../amazonq/commons/types'
 import { ToolkitError } from '../../../shared/errors'
 import * as crypto from '../../../shared/crypto'
 import { createMockSessionStateAction } from '../../amazonq/utils'
 
-import { TestSessionMocks } from '../../amazonq/utils'
-import { beforeEachFunc, createSessionTestSetup } from '../../amazonq/session/testSetup'
-
-let testMocks: TestSessionMocks
+import { createTestContext, setupTestHooks } from '../../amazonq/session/testSetup'
 
 describe('sessionStateFeatureDev', () => {
-    const { conversationId, uploadId, tabId, currentCodeGenerationId } = createSessionTestSetup()
-    let testConfig: SessionStateConfig
-
-    beforeEach(async () => {
-        testMocks = {}
-        testConfig = await beforeEachFunc(testMocks, conversationId, uploadId, currentCodeGenerationId)
-    })
-
-    afterEach(() => {
-        sinon.restore()
-    })
+    const context = createTestContext()
+    setupTestHooks(context)
 
     describe('MockCodeGenState', () => {
         it('loops forever in the same state', async () => {
             sinon.stub(crypto, 'randomUUID').returns('upload-id' as ReturnType<(typeof crypto)['randomUUID']>)
             const testAction = createMockSessionStateAction()
-            const state = new MockCodeGenState(testConfig, tabId)
+            const state = new MockCodeGenState(context.testConfig, context.tabId)
             const result = await state.interact(testAction)
 
             assert.deepStrictEqual(result, {
@@ -51,40 +38,52 @@ describe('sessionStateFeatureDev', () => {
     describe('FeatureDevPrepareCodeGenState', () => {
         it('error when failing to prepare repo information', async () => {
             sinon.stub(vscode.workspace, 'findFiles').throws()
-            testMocks.createUploadUrl!.resolves({ uploadId: '', uploadUrl: '' })
+            context.testMocks.createUploadUrl!.resolves({ uploadId: '', uploadUrl: '' })
             const testAction = createMockSessionStateAction()
 
             await assert.rejects(() => {
-                return new FeatureDevPrepareCodeGenState(testConfig, [], [], [], tabId, 0).interact(testAction)
+                return new FeatureDevPrepareCodeGenState(context.testConfig, [], [], [], context.tabId, 0).interact(
+                    testAction
+                )
             })
         })
     })
 
     describe('FeatureDevCodeGenState', () => {
         it('transitions to FeatureDevPrepareCodeGenState when codeGenerationStatus ready ', async () => {
-            testMocks.getCodeGeneration!.resolves({
+            context.testMocks.getCodeGeneration!.resolves({
                 codeGenerationStatus: { status: 'Complete' },
                 codeGenerationRemainingIterationCount: 2,
                 codeGenerationTotalIterationCount: 3,
             })
 
-            testMocks.exportResultArchive!.resolves({ newFileContents: [], deletedFiles: [], references: [] })
+            context.testMocks.exportResultArchive!.resolves({ newFileContents: [], deletedFiles: [], references: [] })
 
             const testAction = createMockSessionStateAction()
-            const state = new FeatureDevCodeGenState(testConfig, [], [], [], tabId, 0, {}, 2, 3)
+            const state = new FeatureDevCodeGenState(context.testConfig, [], [], [], context.tabId, 0, {}, 2, 3)
             const result = await state.interact(testAction)
 
-            const nextState = new FeatureDevPrepareCodeGenState(testConfig, [], [], [], tabId, 1, 2, 3, undefined)
+            const nextState = new FeatureDevPrepareCodeGenState(
+                context.testConfig,
+                [],
+                [],
+                [],
+                context.tabId,
+                1,
+                2,
+                3,
+                undefined
+            )
 
             assert.deepStrictEqual(result.nextState?.deletedFiles, nextState.deletedFiles)
             assert.deepStrictEqual(result.nextState?.filePaths, result.nextState?.filePaths)
             assert.deepStrictEqual(result.nextState?.references, result.nextState?.references)
         })
 
         it('fails when codeGenerationStatus failed ', async () => {
-            testMocks.getCodeGeneration!.rejects(new ToolkitError('Code generation failed'))
+            context.testMocks.getCodeGeneration!.rejects(new ToolkitError('Code generation failed'))
             const testAction = createMockSessionStateAction()
-            const state = new FeatureDevCodeGenState(testConfig, [], [], [], tabId, 0, {})
+            const state = new FeatureDevCodeGenState(context.testConfig, [], [], [], context.tabId, 0, {})
             try {
                 await state.interact(testAction)
                 assert.fail('failed code generations should throw an error')