From e752a6ac9c336d55f0268f60451c19406f9845e6 Mon Sep 17 00:00:00 2001
From: Gregory Petukhov <lorien@lorien.name>
Date: Fri, 10 Nov 2017 15:28:17 +0600
Subject: [PATCH] Disable token transfers till treasury is set

---
 contracts/TestToken.sol | 14 +++++++++++++-
 test/token.js           | 30 ++++++++++++++++++++++++++++--
 2 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/contracts/TestToken.sol b/contracts/TestToken.sol
index 6025cea..e707028 100644
--- a/contracts/TestToken.sol
+++ b/contracts/TestToken.sol
@@ -23,11 +23,23 @@ contract TestToken is StandardToken, Ownable {
         treasury = Treasury(_treasury);
     }
 
-    function transfer(address _to, uint256 _tokenAmount) public returns (bool) {
+    modifier treasuryIsSet() {
+        require(address(treasury) != 0x0);
+        _;
+    }
+
+    function transfer(address _to, uint256 _tokenAmount) public treasuryIsSet returns (bool) {
         super.transfer(_to, _tokenAmount);
         if (_to == address(treasury)) {
             treasury.tokenDepositEvent(msg.sender, _tokenAmount);
         }
     }
 
+    function transferFrom(address _from, address _to, uint256 _value) public treasuryIsSet returns (bool) {
+        super.transferFrom(_from, _to, _value);
+        if (_to == address(treasury)) {
+            // TODO FIXME
+            //treasury.tokenDepositEvent(msg.sender, _tokenAmount);
+        }
+    }
 }
diff --git a/test/token.js b/test/token.js
index 0b53049..84736ac 100644
--- a/test/token.js
+++ b/test/token.js
@@ -8,7 +8,7 @@ let TestToken = artifacts.require('TestToken');
 let {deployTestContracts} = require('./util/deploy.js');
 
 contract('TestToken [all features]', function(accounts) {
-    let {tokenContract} = {};
+    let {tokenContract, treasuryContract} = {};
 
     beforeEach(async () => {
         ({tokenContract, treasuryContract} = await deployTestContracts(accounts));
@@ -45,6 +45,32 @@ contract('TestToken [all features]', function(accounts) {
         );
     });
 
-});
+    it('transfer fails if treasury is not set', async () => {
+        await tokenContract.transfer(accounts[1], 100)
+            .should.be.rejectedWith('invalid opcode');
+    });
+
+    it('transferFrom fails if treasury is not set', async () => {
+        await tokenContract.approve(accounts[1], 100);
+        await tokenContract.transferFrom(accounts[0], accounts[2], 100, {from: accounts[1]})
+            .should.be.rejectedWith('invalid opcode');
+    });
+
+    it('transfer works if treasury is set', async () => {
+        await tokenContract.setTreasury(treasuryContract.address);
+        await tokenContract.transfer(accounts[1], 100);
+        100..should.be.bignumber.equal(
+            await tokenContract.balanceOf(accounts[1])
+        );
+    });
 
+    it('transferFrom works if treasury is set', async () => {
+        await tokenContract.setTreasury(treasuryContract.address);
+        await tokenContract.approve(accounts[1], 100);
+        await tokenContract.transferFrom(accounts[0], accounts[2], 100, {from: accounts[1]});
+        100..should.be.bignumber.equal(
+            await tokenContract.balanceOf(accounts[2])
+        );
+    });
 
+});