Refactor as per review

pooltogether · Jul 31, 2020 · a34b3ad · a34b3ad
1 parent 5125317
commit a34b3ad
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/contracts/prize-pool/PrizePool.sol b/contracts/prize-pool/PrizePool.sol
@@ -284,9 +284,7 @@ abstract contract PrizePool is OwnableUpgradeSafe, BaseRelayRecipient, Reentranc
   /// @param from The address the tickets are being transferred from
   /// @param to The address the tickets are being transferred to
   /// @param amount The amount of tickets being trasferred
-  function beforeTokenTransfer(address from, address to, uint256 amount) external override {
-    require(isControlled(_msgSender()), "PrizePool/uncontrolled-token");
-
+  function beforeTokenTransfer(address from, address to, uint256 amount) external override onlyControlledToken(msg.sender) {
     // minting and redeeming are handled separately
     if (from != address(0) && to != address(0) && _hasPrizeStrategy()) {
       prizeStrategy.beforeTokenTransfer(from, to, amount, msg.sender);

diff --git a/test/CompoundPrizePool.test.js b/test/CompoundPrizePool.test.js
@@ -60,7 +60,7 @@ describe('CompoundPrizePool', function() {
     describe("beforeTokenTransfer()", () => {
       it('should not allow uncontrolled tokens to call', async () => {
         await expect(prizePool.beforeTokenTransfer(wallet._address, wallet2._address, toWei('1')))
-          .to.be.revertedWith('PrizePool/uncontrolled-token')
+          .to.be.revertedWith('PrizePool/unknown-token')
       })
     })