diff --git a/packages/schematics/angular/refactor/jasmine-vitest/test-file-transformer.ts b/packages/schematics/angular/refactor/jasmine-vitest/test-file-transformer.ts index eb8afdfe0449..79f4ee2cfc9a 100644 --- a/packages/schematics/angular/refactor/jasmine-vitest/test-file-transformer.ts +++ b/packages/schematics/angular/refactor/jasmine-vitest/test-file-transformer.ts @@ -56,6 +56,19 @@ import { RefactorReporter } from './utils/refactor-reporter'; */ const BLANK_LINE_PLACEHOLDER = '// __PRESERVE_BLANK_LINE__'; +/** + * Vitest function names that should be imported when using the --add-imports option. + */ +const VITEST_FUNCTION_NAMES = new Set([ + 'describe', + 'it', + 'expect', + 'beforeEach', + 'afterEach', + 'beforeAll', + 'afterAll', +]); + /** * Replaces blank lines in the content with a placeholder to prevent TypeScript's printer * from removing them. This ensures that the original formatting of blank lines is preserved. @@ -183,7 +196,7 @@ export function transformJasmineToVitest( if (ts.isCallExpression(transformedNode)) { if (options.addImports && ts.isIdentifier(transformedNode.expression)) { const name = transformedNode.expression.text; - if (name === 'describe' || name === 'it' || name === 'expect') { + if (VITEST_FUNCTION_NAMES.has(name)) { addVitestValueImport(pendingVitestValueImports, name); } } diff --git a/packages/schematics/angular/refactor/jasmine-vitest/test-file-transformer_add-imports_spec.ts b/packages/schematics/angular/refactor/jasmine-vitest/test-file-transformer_add-imports_spec.ts index fac88bb9ed02..1fd4beb6546e 100644 --- a/packages/schematics/angular/refactor/jasmine-vitest/test-file-transformer_add-imports_spec.ts +++ b/packages/schematics/angular/refactor/jasmine-vitest/test-file-transformer_add-imports_spec.ts @@ -81,5 +81,41 @@ describe('Jasmine to Vitest Transformer', () => { `; await expectTransformation(input, expected, true); }); + + it('should add imports for beforeEach and afterEach when addImports is true', async () => { + const input = ` + describe('My Suite', () => { + beforeEach(() => {}); + afterEach(() => {}); + }); + `; + const expected = ` + import { afterEach, beforeEach, describe } from 'vitest'; + + describe('My Suite', () => { + beforeEach(() => {}); + afterEach(() => {}); + }); + `; + await expectTransformation(input, expected, true); + }); + + it('should add imports for beforeAll and afterAll when addImports is true', async () => { + const input = ` + describe('My Suite', () => { + beforeAll(() => {}); + afterAll(() => {}); + }); + `; + const expected = ` + import { afterAll, beforeAll, describe } from 'vitest'; + + describe('My Suite', () => { + beforeAll(() => {}); + afterAll(() => {}); + }); + `; + await expectTransformation(input, expected, true); + }); }); });