fix: refactor remappings and add test import

src/smock-foundry.ts

@@ -8,6 +8,8 @@ import {
   smockableNode,
   compileSolidityFilesFoundry,
   renderAbstractUnimplementedFunctions,
+  getRemappings,
+  getTestImport,
 } from './utils';
 import path from 'path';
 import { ensureDir } from 'fs-extra';
@@ -31,7 +33,9 @@ export async function generateMockContracts(
     console.log('Parsing contracts...');
 
     try {
-      const sourceUnits = await getSourceUnits(rootPath, contractsDirectories, ignoreDirectories);
+      const remappings: string[] = await getRemappings(rootPath);
+      const testImport = getTestImport(remappings);
+      const sourceUnits = await getSourceUnits(rootPath, contractsDirectories, ignoreDirectories, remappings);
 
       if (!sourceUnits.length) return console.error('No solidity files found in the specified directory');
 
@@ -80,6 +84,7 @@ export async function generateMockContracts(
             sourceContractRelativePath: sourceContractRelativePath,
             exportedSymbols: Array.from(scope.exportedSymbols.keys()),
             license: sourceUnit.license,
+            testImport,
           });
 
           await ensureDir(path.dirname(mockContractAbsolutePath));
@@ -89,12 +94,12 @@ export async function generateMockContracts(
 
       // Generate SmockHelper contract
       const smockHelperTemplate = await getSmockHelperTemplate();
-      const smockHelperCode: string = smockHelperTemplate({});
+      const smockHelperCode: string = smockHelperTemplate({ testImport });
       writeFileSync(`${mocksPath}/SmockHelper.sol`, smockHelperCode);
 
       console.log('Mock contracts generated successfully');
 
-      await compileSolidityFilesFoundry(rootPath, mocksDirectory);
+      await compileSolidityFilesFoundry(rootPath, mocksDirectory, remappings);
 
       // Format the generated files
       console.log('Formatting generated files...');

src/templates/contract-template.hbs

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: {{#if license}}{{license}}{{else}}UNLICENSED{{/if}}
 pragma solidity ^0.8.0;
 
-import { Test } from 'forge-std/Test.sol';
+import { Test } from '{{testImport}}';
 import { {{~exportedSymbols~}} } from '{{sourceContractRelativePath}}';
 {{importsContent}}
 

src/templates/helper-template.hbs

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.0;
 
-import {Test} from 'forge-std/Test.sol';
+import {Test} from '{{testImport}}';
 
 contract SmockHelper is Test {
   function deployMock(

src/utils.ts

@@ -31,6 +31,7 @@
   externalOrPublicFunctionContext,
   internalFunctionContext,
 } from './context';
+import { exec } from 'child_process';
 
 /**
  * Fixes user-defined types
@@ -73,7 +74,7 @@
  * @param templatePath The path of the template (if it's nested)
  * @returns The compiled template
  */
 export function compileTemplate(templateName: string, templatePath?: string[]): HandlebarsTemplateDelegate<any> {
   const templateContent = readTemplate(templateName, templatePath);
   return Handlebars.compile(templateContent, { noEscape: true });
 }
@@ -82,7 +83,7 @@
  * Gets the base contract template
  * @returns The contract template
  */
 export function getContractTemplate(): HandlebarsTemplateDelegate<any> {
   return compileTemplate('contract-template');
 }
 
@@ -90,7 +91,7 @@
  * Gets the smock helper template
  * @returns The helper template
  */
 export function getSmockHelperTemplate(): HandlebarsTemplateDelegate<any> {
   return compileTemplate('helper-template');
 }
 
@@ -98,13 +99,11 @@
  * Compiles the solidity files in the given directory calling forge build command
  * @param mockContractsDir The directory of the generated contracts
  */
-export async function compileSolidityFilesFoundry(rootPath: string, mockContractsDir: string) {
+export async function compileSolidityFilesFoundry(rootPath: string, mockContractsDir: string, remappings: string[]) {
   console.log('Compiling contracts...');
   try {
     const solidityFiles: string[] = await getSolidityFilesAbsolutePaths(rootPath, [mockContractsDir]);
 
-    const remappings: string[] = await getRemappings(rootPath);
-
     await compileSol(solidityFiles, 'auto', {
       basePath: rootPath,
       remapping: remappings,
@@ -245,7 +244,12 @@
     try {
       return await exports.getRemappingsFromConfig(path.join(rootPath, 'foundry.toml'));
     } catch {
-      return [];
+      // If neither file exists, try to generate the remappings using forge
+      try {
+        return await exports.getRemappingsFromForge();
+      } catch {
+        return [];
+      }
     }
   }
 }
@@ -265,17 +269,32 @@
   const regex = /remappings[\s|\n]*=[\s\n]*\[(?<remappings>[^\]]+)]/;
   const matches = foundryConfigContent.match(regex);
   if (matches) {
     return matches
       .groups!.remappings.split(',')
       .map((line) => line.trim())
       .map((line) => line.replace(/["']/g, ''))
       .filter((line) => line.length)
       .map((line) => sanitizeRemapping(line));
   } else {
-    return [];
+    throw new Error('No remappings found in foundry.toml');
   }
 }
 
+/**
+ * Returns the remappings generated by forge
+ * @returns {Promise<string[]>} - The list of remappings
+ */
+export async function getRemappingsFromForge(): Promise<string[]> {
+  const remappingsContent = await new Promise<string>((resolve, reject) =>
+    exec('forge remappings', { encoding: 'utf8' }, (error, stdout) => (error ? reject(error) : resolve(stdout))),
+  );
+  return remappingsContent
+    .split('\n')
+    .map((line) => line.trim())
+    .filter((line) => line.length)
+    .map((line) => sanitizeRemapping(line));
+}
+
 export function sanitizeRemapping(line: string): string {
   // Make sure the key and the value both either have or don't have a trailing slash
   const [key, value] = line.split('=');
@@ -304,12 +323,30 @@
   }
 }
 
-export async function getSourceUnits(rootPath: string, contractsDirectories: string[], ignoreDirectories: string[]): Promise<SourceUnit[]> {
+export function getTestImport(remappings: string[]): string {
+  const module = 'forge-std';
+
+  for (const remapping of remappings) {
+    const [alias, path] = remapping.split('='); // Split remapping into alias and path
+
+    if (alias.startsWith(module) && path.includes(module)) {
+      const srcPath = path.includes('/src/') ? '' : `src/`;
+      return `${alias}${srcPath}Test.sol`;
+    }
+  }
+
+  return 'forge-std/src/Test.sol';
+}
+
+export async function getSourceUnits(
+  rootPath: string,
+  contractsDirectories: string[],
+  ignoreDirectories: string[],
+  remappings: string[],
+): Promise<SourceUnit[]> {
   const files: string[] = await getSolidityFilesAbsolutePaths(rootPath, contractsDirectories);
   const solidityFiles = files.filter((file) => !ignoreDirectories.some((directory) => file.includes(directory)));
 
-  const remappings: string[] = await getRemappings(rootPath);
-
   const compiledFiles = await compileSol(solidityFiles, 'auto', {
     basePath: rootPath,
     remapping: remappings,