add tests for migrate-imports.js and fix for upgradeable

Author: frangio

scripts/migrate-imports.js

@@ -3,11 +3,6 @@
 const { promises: fs } = require('fs');
 const path = require('path');
 
-const versions = [
-  '@openzeppelin/contracts',
-  '@openzeppelin/contracts-upgradeable',
-];
-
 const pathUpdates = {
   // 'access/AccessControl.sol': undefined,
   // 'access/Ownable.sol': undefined,
@@ -112,15 +107,15 @@ async function main (paths = [ 'contracts' ]) {
   }
 }
 
-async function listFilesRecursively (paths, filter = undefined) {
-  const queue = Array.isArray(paths) ? paths : [ paths ];
+async function listFilesRecursively (paths, filter) {
+  const queue = paths;
   const files = [];
 
   while (queue.length > 0) {
     const top = queue.shift();
     const stat = await fs.stat(top);
     if (stat.isFile()) {
-      if (!filter || top.match(filter)) {
+      if (top.match(filter)) {
         files.push(top);
       }
     } else if (stat.isDirectory()) {
@@ -146,18 +141,29 @@ async function updateFile (file, update) {
 
 function updateImportPaths (source) {
   for (const [ oldPath, newPath ] of Object.entries(pathUpdates)) {
-    for (const version of versions) {
-      source = source.replace(
-        path.join(version, oldPath),
-        path.join(version, newPath),
-      );
-    }
+    source = source.replace(
+      path.join('@openzeppelin/contracts', oldPath),
+      path.join('@openzeppelin/contracts', newPath),
+    );
+    source = source.replace(
+      path.join('@openzeppelin/contracts-upgradeable', getUpgradeablePath(oldPath)),
+      path.join('@openzeppelin/contracts-upgradeable', getUpgradeablePath(newPath)),
+    );
   }
 
   return source;
 }
 
-module.exports = main;
+function getUpgradeablePath (file) {
+  const { dir, name, ext } = path.parse(file);
+  const upgradeableName = name + 'Upgradeable';
+  return path.format({ dir, ext, name: upgradeableName });
+}
+
+module.exports = {
+  pathUpdates,
+  updateImportPaths,
+};
 
 if (require.main === module) {
   const args = process.argv.length > 2 ? process.argv.slice(2) : undefined;

test/migrate-imports.test.js

@@ -0,0 +1,25 @@
+const path = require('path');
+const { promises: fs, constants: { F_OK } } = require('fs');
+const { expect } = require('chai');
+
+const { pathUpdates, updateImportPaths } = require('../scripts/migrate-imports.js');
+
+describe('migrate-imports.js', function () {
+  it('every new path exists', async function () {
+    for (const p of Object.values(pathUpdates)) {
+      await fs.access(path.join('contracts', p), F_OK);
+    }
+  });
+
+  it('replaces import paths in a file', async function () {
+    const source = `
+import '@openzeppelin/contracts/math/Math.sol';
+import '@openzeppelin/contracts-upgradeable/math/MathUpgradeable.sol';
+    `;
+    const expected = `
+import '@openzeppelin/contracts/utils/math/Math.sol';
+import '@openzeppelin/contracts-upgradeable/utils/math/MathUpgradeable.sol';
+    `;
+    expect(updateImportPaths(source)).to.equal(expected);
+  });
+});