diff --git a/contracts/Staking.sol b/contracts/Staking.sol
index f81b76a..c9984de 100644
--- a/contracts/Staking.sol
+++ b/contracts/Staking.sol
@@ -19,6 +19,7 @@ contract Staking is Autopetrified, ERCStaking, ERCStakingHistory, IStakingLockin
     string private constant ERROR_TOKEN_NOT_CONTRACT = "STAKING_TOKEN_NOT_CONTRACT";
     string private constant ERROR_AMOUNT_ZERO = "STAKING_AMOUNT_ZERO";
     string private constant ERROR_TOKEN_TRANSFER = "STAKING_TOKEN_TRANSFER";
+    string private constant ERROR_TOKEN_DEPOSIT = "STAKING_TOKEN_DEPOSIT";
     string private constant ERROR_TOKEN_NOT_SENDER = "STAKING_TOKEN_NOT_SENDER";
     string private constant ERROR_WRONG_TOKEN = "STAKING_WRONG_TOKEN";
     string private constant ERROR_NOT_ENOUGH_BALANCE = "STAKING_NOT_ENOUGH_BALANCE";
@@ -441,7 +442,7 @@ contract Staking is Autopetrified, ERCStaking, ERCStakingHistory, IStakingLockin
         _modifyTotalStaked(_amount, true);
 
         // pull tokens into Staking contract
-        require(stakingToken.safeTransferFrom(_from, this, _amount), ERROR_TOKEN_TRANSFER);
+        require(stakingToken.safeTransferFrom(_from, this, _amount), ERROR_TOKEN_DEPOSIT);
 
         emit Staked(_accountAddress, _amount, newStake, _data);
     }
@@ -453,7 +454,8 @@ contract Staking is Autopetrified, ERCStaking, ERCStakingHistory, IStakingLockin
         if (_increase) {
             newStake = currentStake.add(_by);
         } else {
-            newStake = currentStake.sub(_by);
+            require(currentStake >= _by, ERROR_NOT_ENOUGH_BALANCE);
+            newStake = currentStake - _by;
         }
 
         // add new value to account history
diff --git a/package.json b/package.json
index 4a8aad4..2f02082 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@aragon/staking",
-  "version": "0.1.0",
+  "version": "0.1.1",
   "description": "",
   "author": "Aragon Institution MTU <contact@aragon.one>",
   "license": "(GPL-3.0-or-later OR MIT)",
diff --git a/test/helpers/errors.js b/test/helpers/errors.js
index 53769b6..5732712 100644
--- a/test/helpers/errors.js
+++ b/test/helpers/errors.js
@@ -12,6 +12,7 @@ const STAKING_ERRORS = {
   ERROR_TOKEN_NOT_CONTRACT: 'STAKING_TOKEN_NOT_CONTRACT',
   ERROR_AMOUNT_ZERO: 'STAKING_AMOUNT_ZERO',
   ERROR_TOKEN_TRANSFER: 'STAKING_TOKEN_TRANSFER',
+  ERROR_TOKEN_DEPOSIT: 'STAKING_TOKEN_DEPOSIT',
   ERROR_TOKEN_NOT_SENDER: 'STAKING_TOKEN_NOT_SENDER',
   ERROR_WRONG_TOKEN: 'STAKING_WRONG_TOKEN',
   ERROR_NOT_ENOUGH_BALANCE: 'STAKING_NOT_ENOUGH_BALANCE',
diff --git a/test/staking.js b/test/staking.js
index 1155fcc..fc4849e 100644
--- a/test/staking.js
+++ b/test/staking.js
@@ -3,7 +3,7 @@ const { bn, assertBn } = require('@aragon/contract-helpers-test/numbers')
 
 const { deploy } = require('./helpers/deploy')(artifacts)
 const { DEFAULT_STAKE_AMOUNT, EMPTY_DATA } = require('./helpers/constants')
-const { STAKING_ERRORS, SAFE_MATH_ERRORS } = require('./helpers/errors')
+const { STAKING_ERRORS } = require('./helpers/errors')
 
 const StakingMock = artifacts.require('StakingMock')
 const StandardTokenMock = artifacts.require('StandardTokenMock');
@@ -66,7 +66,7 @@ contract('Staking app', ([owner, other]) => {
     const balance = await getTokenBalance(token, owner)
     const amount = balance + 1
     await token.approve(stakingAddress, amount)
-    await assertRevert(staking.stake(amount, EMPTY_DATA), STAKING_ERRORS.ERROR_TOKEN_TRANSFER)
+    await assertRevert(staking.stake(amount, EMPTY_DATA), STAKING_ERRORS.ERROR_TOKEN_DEPOSIT)
   })
 
   it('stakes for', async () => {
@@ -112,7 +112,7 @@ contract('Staking app', ([owner, other]) => {
 
   it('fails unstaking more than staked', async () => {
     await approveAndStake()
-    await assertRevert(staking.unstake(DEFAULT_STAKE_AMOUNT + 1, EMPTY_DATA), SAFE_MATH_ERRORS.ERROR_SUB_UNDERFLOW)
+    await assertRevert(staking.unstake(DEFAULT_STAKE_AMOUNT + 1, EMPTY_DATA), STAKING_ERRORS.ERROR_NOT_ENOUGH_BALANCE)
   })
 
   context('History', async () => {