Merge pull request #5 from vdrg/puzzle-wallet-tests

OpenZeppelin · Oct 7, 2021 · 98649bc · 98649bc
2 parents 67f568c + 3ec9dfa
commit 98649bc
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 16 deletions.
diff --git a/contracts/contracts/levels/PuzzleWallet.sol b/contracts/contracts/levels/PuzzleWallet.sol
@@ -55,11 +55,6 @@ contract PuzzleWallet {
         _;
     }
 
-    // TODO: remove thisssss
-    function takeOwnership() external {
-        owner = msg.sender;
-    }
-
     function addToWhitelist(address addr) external onlyOwner {
         whitelisted[addr] = true;
     }

diff --git a/contracts/contracts/levels/PuzzleWalletFactory.sol b/contracts/contracts/levels/PuzzleWalletFactory.sol
@@ -8,16 +8,24 @@ contract PuzzleWalletFactory is Level {
   function createInstance(address /*_player*/) override public payable returns (address) {
     require(msg.value == 1 ether, "Must send 1 ETH to create instance");
 
-    PuzzleWallet puzzle = new PuzzleWallet();
+    // Deploy the Wallet logic
+    PuzzleWallet walletLogic = new PuzzleWallet();
+    walletLogic.init();
 
-    puzzle.addToWhitelist(address(this));
+    // Proxy instance
+    bytes memory data= abi.encodeWithSelector(PuzzleWallet.init.selector, 100 ether);
+    PuzzleProxy proxy = new PuzzleProxy(address(this), address(walletLogic), data);
 
-    puzzle.deposit{value: 1 ether}(1 ether);
+    PuzzleWallet instance = PuzzleWallet(address(proxy));
 
-    return address(puzzle);
+    instance.addToWhitelist(address(this));
+    instance.deposit{ value: msg.value }(1 ether);
+
+    return address(proxy);
   }
 
-  function validateInstance(address payable _instance, address /*_player*/) override public returns (bool) {
-    return _instance.balance == 0;
+  function validateInstance(address payable _instance, address _player) override public returns (bool) {
+    PuzzleWallet instance = PuzzleWallet(_instance);
+    return _instance.balance == 0 && instance.owner() == _player && instance.whitelisted(_player);
   }
 }
diff --git a/contracts/test/levels/PuzzleWallet.test.js b/contracts/test/levels/PuzzleWallet.test.js
@@ -1,8 +1,10 @@
 const ethutil = require('ethereumjs-util')
-const PuzzleWalletFactory = artifacts.require('./levels/PuzzleWalletFactory.sol')
-const PuzzleWallet = artifacts.require('./attacks/PuzzleWallet.sol')
 
+const PuzzleProxy  = artifacts.require('PuzzleProxy');
+const PuzzleWalletFactory = artifacts.require('PuzzleWalletFactory')
+const PuzzleWallet = artifacts.require('PuzzleWallet')
 const Ethernaut = artifacts.require('./Ethernaut.sol')
+
 const { BN, constants, expectEvent, expectRevert } = require('openzeppelin-test-helpers')
 const utils = require('../utils/TestUtils')
 
@@ -26,11 +28,17 @@ contract('PuzzleWallet', function(accounts) {
       {from: player, value: web3.utils.toWei('1', 'ether')}
     )
 
-    assert.equal(await instance.owner(), level.address)
+    assert.equal(level.address, await instance.owner(), 'Owner is not the factory')
     assert.equal(web3.utils.toWei('1', 'ether'), (await instance.balances(level.address)).toString())
+
+    const proxy = await PuzzleProxy.at(instance.address)
+    await proxy.proposeNewAdmin(player)
+
+    // check that the player has placed their address in the owner slot
+    assert.equal(player, await instance.owner(), "Player is not the owner")
 
-    // TODO: take ownership in a real way
-    await instance.takeOwnership({from: player})
+    // check that player is not whitelisted yet
+    assert.isFalse(await instance.whitelisted(player), 'Player is not whitelisted')
 
     // Player whitelists herself
     await instance.addToWhitelist(player, { from: player })