fix(runner): preserve handler wrapping on extend (#8153)

Author: pengooseDev

packages/runner/src/suite.ts

@@ -770,7 +770,8 @@ export function createTaskCollector(
       runner,
     )
 
-    return createTest(function fn(
+    const originalWrapper = fn
+    return createTest(function (
       name: string | Function,
       optionsOrFn?: TestOptions | TestFunction,
       optionsOrTest?: number | TestOptions | TestFunction,
@@ -784,12 +785,9 @@ export function createTaskCollector(
           scopedFixtures,
         )
       }
-      collector.test.fn.call(
-        context,
-        formatName(name),
-        optionsOrFn as TestOptions,
-        optionsOrTest as TestFunction,
-      )
+      const { handler, options } = parseArguments(optionsOrFn, optionsOrTest)
+      const timeout = options.timeout ?? runner?.config.testTimeout
+      originalWrapper.call(context, formatName(name), handler, timeout)
     }, _context)
   }
 

test/core/test/task-collector.test.ts

@@ -1,5 +1,5 @@
-import { expect, test, vi } from 'vitest'
-import { createTaskCollector } from 'vitest/suite'
+import { assert, describe, expect, test, vi } from 'vitest'
+import { createTaskCollector, getCurrentSuite } from 'vitest/suite'
 
 test('collector keeps the order of arguments', () => {
   const fn = vi.fn()
@@ -23,3 +23,29 @@ test('collector keeps the order of arguments', () => {
 
   expect(fn).toHaveBeenNthCalledWith(4, 'a', options, expect.any(Function))
 })
+
+describe('collector.extend should preserve handler wrapping', () => {
+  let flag = false
+
+  const flagTest = createTaskCollector(function (
+    this: object,
+    name: string,
+    fn: () => void,
+  ) {
+    const handler = async () => {
+      flag = false
+      await fn()
+      assert(flag)
+    }
+    getCurrentSuite().task(name, { ...this, handler })
+  })
+
+  const extendedTest = flagTest.extend({})
+
+  extendedTest.fails('should fail when flag is never set', {}, () => {})
+
+  flagTest('should pass when flag is set', () => {
+    flag = true
+    expect(flag).toBe(true)
+  })
+})