From 3772233cf5c07c33c3830aece22419ee5dd353ed Mon Sep 17 00:00:00 2001
From: nikeshnazareth <business@nikeshnazareth.com>
Date: Sun, 3 Mar 2019 02:36:36 +1100
Subject: [PATCH] Add guard to ERC20Migrator migrate function (#1659)

* Add guard to ERC20Migrator migrate function

* Add tests for premature migration in ERC20Migrator

These tests apply to the new guard condition, but they don't
fail without it, since the functions revert anyway.
They are added for completeness and to ensure full code coverage.

* Use context block around premature migration tests

We should use context blocks for situational details
and describe for features or functions.
---
 contracts/drafts/ERC20Migrator.sol |  1 +
 test/drafts/ERC20Migrator.test.js  | 34 ++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/contracts/drafts/ERC20Migrator.sol b/contracts/drafts/ERC20Migrator.sol
index 3a311383a..18ad0df3a 100644
--- a/contracts/drafts/ERC20Migrator.sol
+++ b/contracts/drafts/ERC20Migrator.sol
@@ -82,6 +82,7 @@ contract ERC20Migrator {
      * @param amount amount of tokens to be migrated
      */
     function migrate(address account, uint256 amount) public {
+        require(address(_newToken) != address(0));
         _legacyToken.safeTransferFrom(account, address(this), amount);
         _newToken.mint(account, amount);
     }
diff --git a/test/drafts/ERC20Migrator.test.js b/test/drafts/ERC20Migrator.test.js
index bc2a6e058..6e8fd451d 100644
--- a/test/drafts/ERC20Migrator.test.js
+++ b/test/drafts/ERC20Migrator.test.js
@@ -44,6 +44,40 @@ contract('ERC20Migrator', function ([_, owner, recipient, anotherAccount]) {
       });
     });
 
+    context('before starting the migration', function () {
+      it('returns the zero address for the new token', async function () {
+        (await this.migrator.newToken()).should.be.equal(ZERO_ADDRESS);
+      });
+
+      describe('migrateAll', function () {
+        const amount = totalSupply;
+
+        describe('when the approved balance is equal to the owned balance', function () {
+          beforeEach('approving the whole balance to the new contract', async function () {
+            await this.legacyToken.approve(this.migrator.address, amount, { from: owner });
+          });
+
+          it('reverts', async function () {
+            await shouldFail.reverting(this.migrator.migrateAll(owner));
+          });
+        });
+      });
+
+      describe('migrate', function () {
+        const amount = new BN(50);
+
+        describe('when the amount is equal to the approved value', function () {
+          beforeEach('approving tokens to the new contract', async function () {
+            await this.legacyToken.approve(this.migrator.address, amount, { from: owner });
+          });
+
+          it('reverts', async function () {
+            await shouldFail.reverting(this.migrator.migrate(owner, amount));
+          });
+        });
+      });
+    });
+
     describe('once migration began', function () {
       beforeEach('beginning migration', async function () {
         await this.newToken.addMinter(this.migrator.address);