From 69ed7235ab98279f673d861598aa1f36fece0648 Mon Sep 17 00:00:00 2001
From: Aapo Kiiso <aapo@kii.so>
Date: Mon, 17 Jul 2023 14:42:05 +0100
Subject: [PATCH] feat: register all auto-imports for mocking (#254)

Co-authored-by: Ghazi <ghazialhouwari@gmail.com>
---
 src/vitest-environment-nuxt/modules/mock.ts     | 17 ++++++-----------
 test/fixtures/nuxt-vitest/modules/custom.ts     | 17 +++++++++++++++--
 .../runtime/composables/auto-import-mock.ts     |  7 +++++++
 .../tests/nuxt/auto-import-mock.spec.ts         |  9 +++++++++
 4 files changed, 37 insertions(+), 13 deletions(-)
 create mode 100644 test/fixtures/nuxt-vitest/modules/runtime/composables/auto-import-mock.ts

diff --git a/src/vitest-environment-nuxt/modules/mock.ts b/src/vitest-environment-nuxt/modules/mock.ts
index 88f2fa41d..f322dfd4c 100644
--- a/src/vitest-environment-nuxt/modules/mock.ts
+++ b/src/vitest-environment-nuxt/modules/mock.ts
@@ -1,4 +1,4 @@
-import type { Import } from 'unimport'
+import type { Import, Unimport } from 'unimport'
 import { addVitePlugin, defineNuxtModule } from '@nuxt/kit'
 import { walk } from 'estree-walker'
 import type { CallExpression } from 'estree'
@@ -25,8 +25,6 @@ export interface MockComponentInfo {
   factory: string
 }
 
-const nuxtImportSources = ['#app', '#vue-router', 'vue-demi', '@unhead/vue']
-
 /**
  * This module is a macro that transforms `mockNuxtImport()` to `vi.mock()`,
  * which make it possible to mock Nuxt imports.
@@ -36,22 +34,19 @@ export default defineNuxtModule({
     name: PLUGIN_NAME,
   },
   setup(_, nuxt) {
+    let importsCtx: Unimport
     let imports: Import[] = []
     let components: Component[] = []
 
-    nuxt.hook('imports:extend', _ => {
-      imports = imports.concat(_)
-    })
     nuxt.hook('imports:context', async ctx => {
-      // add core nuxt composables to imports
-      const registeredImports = await ctx.getImports()
-      imports = imports.concat(
-        registeredImports.filter(item => nuxtImportSources.includes(item.from))
-      )
+      importsCtx = ctx
     })
     nuxt.hook('components:extend', _ => {
       components = _
     })
+    nuxt.hook('ready', async () => {
+      imports = await importsCtx.getImports()
+    })
 
     // Polyfill Array.prototype.findLastIndex for legacy Node.js
     function findLastIndex<T>(arr: T[], predicate: (item: T) => boolean) {
diff --git a/test/fixtures/nuxt-vitest/modules/custom.ts b/test/fixtures/nuxt-vitest/modules/custom.ts
index f56794af1..b60d6fd19 100644
--- a/test/fixtures/nuxt-vitest/modules/custom.ts
+++ b/test/fixtures/nuxt-vitest/modules/custom.ts
@@ -1,4 +1,4 @@
-import { defineNuxtModule } from '@nuxt/kit'
+import { defineNuxtModule, createResolver, addImports } from '@nuxt/kit'
 
 export default defineNuxtModule({
   meta: {
@@ -6,5 +6,18 @@ export default defineNuxtModule({
   },
   setup(_, _nuxt) {
     console.log('From custom module!')
-  },
+
+    const { resolve } = createResolver(import.meta.url)
+
+    addImports([
+      {
+        name: 'useCustomModuleAutoImportedTarget',
+        from: resolve('runtime/composables/auto-import-mock')
+      },
+      {
+        name: 'useCustomModuleAutoImportedNonTarget',
+        from: resolve('runtime/composables/auto-import-mock')
+      }
+    ])
+  }
 })
diff --git a/test/fixtures/nuxt-vitest/modules/runtime/composables/auto-import-mock.ts b/test/fixtures/nuxt-vitest/modules/runtime/composables/auto-import-mock.ts
new file mode 100644
index 000000000..827066ba5
--- /dev/null
+++ b/test/fixtures/nuxt-vitest/modules/runtime/composables/auto-import-mock.ts
@@ -0,0 +1,7 @@
+export function useCustomModuleAutoImportedTarget() {
+  return 'the original'
+}
+
+export function useCustomModuleAutoImportedNonTarget() {
+  return 'the original'
+}
diff --git a/test/fixtures/nuxt-vitest/tests/nuxt/auto-import-mock.spec.ts b/test/fixtures/nuxt-vitest/tests/nuxt/auto-import-mock.spec.ts
index 9c5252bd3..72e92643b 100644
--- a/test/fixtures/nuxt-vitest/tests/nuxt/auto-import-mock.spec.ts
+++ b/test/fixtures/nuxt-vitest/tests/nuxt/auto-import-mock.spec.ts
@@ -5,8 +5,17 @@ mockNuxtImport<typeof useAutoImportedTarget>('useAutoImportedTarget', () => {
   return () => 'mocked!'
 })
 
+mockNuxtImport<typeof useCustomModuleAutoImportedTarget>('useCustomModuleAutoImportedTarget', () => {
+  return () => 'mocked!'
+})
+
 it('should mock', () => {
   vi.fn()
   expect(useAutoImportedTarget()).toMatchInlineSnapshot('"mocked!"')
   expect(useAutoImportedNonTarget()).toMatchInlineSnapshot('"the original"')
 })
+
+it('should mock composable from external package', () => {
+  expect(useCustomModuleAutoImportedTarget()).toMatchInlineSnapshot('"mocked!"')
+  expect(useCustomModuleAutoImportedNonTarget()).toMatchInlineSnapshot('"the original"')
+})