diff --git a/packages/@sanity/migrate/src/runner/__tests__/normalizeMigrationDefinition.test.ts b/packages/@sanity/migrate/src/runner/__tests__/normalizeMigrationDefinition.test.ts
index 3fbfda30b69..15578ab1c82 100644
--- a/packages/@sanity/migrate/src/runner/__tests__/normalizeMigrationDefinition.test.ts
+++ b/packages/@sanity/migrate/src/runner/__tests__/normalizeMigrationDefinition.test.ts
@@ -1,26 +1,32 @@
 import {describe, expect, it, jest} from '@jest/globals'
+import {type SanityDocument} from '@sanity/types'
 
 import {createIfNotExists} from '../../mutations'
-import {type Migration, type NodeMigration} from '../../types'
+import {
+  type DocumentMigrationReturnValue,
+  type Migration,
+  type MigrationContext,
+  type NodeMigration,
+} from '../../types'
 import {
   createAsyncIterableMutation,
   normalizeMigrateDefinition,
 } from '../normalizeMigrateDefinition'
 
 const mockAsyncIterableIterator = () => {
-  const data = [{_id: 'mockId', _type: 'mockDocumentType'}]
-  let index = 0
-
-  return {
-    next: jest.fn().mockImplementation(() => {
-      if (index < data.length) {
-        return Promise.resolve({value: data[index++], done: false})
-      }
-      return Promise.resolve({value: undefined, done: true})
-    }),
-    [Symbol.asyncIterator]: jest.fn().mockImplementation(function (this: unknown) {
-      return this
-    }),
+  const data: SanityDocument[] = [
+    {
+      _id: 'mockId',
+      _type: 'mockDocumentType',
+      _updatedAt: '2024-02-16T14:13:59Z',
+      _rev: 'xyz',
+      _createdAt: '2024-02-16T14:13:59Z',
+    },
+  ]
+  return async function* documents() {
+    for (let index = 0; index < data.length; index++) {
+      yield data[index]
+    }
   }
 }
 
@@ -58,7 +64,7 @@ describe('#normalizeMigrateDefinition', () => {
     const result = normalizeMigrateDefinition(mockMigration)
     const res = []
 
-    for await (const item of result(mockAsyncIterableIterator as any, {} as any)) {
+    for await (const item of result(mockAsyncIterableIterator(), {} as any)) {
       res.push(item)
     }
 
@@ -79,7 +85,7 @@ describe('#normalizeMigrateDefinition', () => {
     const result = normalizeMigrateDefinition(mockMigration)
     const res = []
 
-    for await (const item of result(mockAsyncIterableIterator as any, {} as any)) {
+    for await (const item of result(mockAsyncIterableIterator(), {} as any)) {
       res.push(item)
     }
 
@@ -90,14 +96,14 @@ describe('#normalizeMigrateDefinition', () => {
 describe('#createAsyncIterableMutation', () => {
   it('should return an async iterable', async () => {
     const mockMigration: NodeMigration = {
-      document: jest.fn() as any,
+      document: jest.fn<() => DocumentMigrationReturnValue>(),
     }
 
     const iterable = createAsyncIterableMutation(mockMigration, {documentTypes: ['foo']})
 
     expect(typeof iterable).toBe('function')
 
-    const iterator = iterable(mockAsyncIterableIterator() as any, {} as any)
+    const iterator = iterable(mockAsyncIterableIterator(), {} as MigrationContext)
     expect(typeof iterator.next).toBe('function')
     expect(typeof iterator.return).toBe('function')
     expect(typeof iterator.throw).toBe('function')