Added deposit/withdraw tests

contracts/BLSWallet.sol

@@ -9,40 +9,50 @@ import "hardhat/console.sol";
 
 contract BLSWallet //is IERC20 //(to consider?)
 {
-  IERC20 baseToken;
+    IERC20 baseToken;
 
-  mapping (address => uint256[4])public blsKeys;
-  mapping (address => uint256) balances;
+    mapping (address => uint256[4]) blsKeys;
+    mapping (address => uint256) balances;
 
-  constructor(IERC20 token) {
-    baseToken = token;
-  }
+    constructor(IERC20 token) {
+        baseToken = token;
+    }
 
-  /**
-  @dev Called from token holder's address
-   */
-  function deposit(
-    uint256[4] memory blsPubKey,
-    uint256 amount
-  ) public {
-    // TODO: check existing key
-    baseToken.transferFrom(msg.sender, address(this), amount);
-    blsKeys[msg.sender] = blsPubKey;
-    balances[msg.sender] += amount;
-  }
+    function senderKeyExists() public pure returns (bool) {
+    }
 
-  function withdraw() public {
-    uint256 amount = balances[msg.sender];
-    blsKeys[msg.sender] = [0,0,0,0];
-    balances[msg.sender] = 0;
-    baseToken.transfer(msg.sender, amount);
-  }
+    /**
+    Begin use of wallet by locking up tokens.
+    @dev Called from token holder's address. Tokens pre-approved.
+    @param blsPubKey BLS public key to be used with address
+    @param amount Amount of tokens the wallet is to take
+     */
+    function deposit(
+        uint256[4] memory blsPubKey,
+        uint256 amount
+    ) public {
+        // TODO: check existing key
+        baseToken.transferFrom(msg.sender, address(this), amount);
+        blsKeys[msg.sender] = blsPubKey;
+        balances[msg.sender] += amount;
+    }
 
-  //   //TODO: verifyMultiple
-  // function transferBatch() {}
+    /**
+    Finish using wallet by returning balance and resetting key.
+    @dev Called from token holder's address
+    */
+    function withdraw() public {
+        uint256 amount = balances[msg.sender];
+        blsKeys[msg.sender] = [0,0,0,0];
+        balances[msg.sender] = 0;
+        baseToken.transfer(msg.sender, amount);
+    }
 
-  function balanceOf(address account) public view returns (uint256) {
-    return balances[account];
-  }
+    //   //TODO: verifyMultiple
+    // function transferBatch() {}
+
+    function balanceOf(address account) public view returns (uint256) {
+        return balances[account];
+    }
 
 }

contracts/mock/MockBLSWallet.sol

@@ -0,0 +1,18 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+import "../BLSWallet.sol";
+
+contract MockBLSWallet is BLSWallet {
+    constructor(
+        IERC20 token
+    ) BLSWallet(token) {
+        
+    }
+
+    function blsPubKeyOf(
+        address account
+    ) public view returns (uint256[4] memory) {
+        return blsKeys[account];
+    }
+}

test/blsWallet-test.js

@@ -1,22 +1,54 @@
 const { expect } = require("chai");
+const { BigNumber } = require("ethers");
 const { ethers } = require("hardhat");
+const utils = ethers.utils;
+
+const acc1BLSPubKey = [0xAA, 0xBB, 0xCC, 0xDD].map((n) => { return BigNumber.from(n) });
+const acc2BLSPubKey = [0xEE, 0xFF, 0x00, 0x11].map((n) => { return BigNumber.from(n) });
+const zeroBLSPubKey = [0, 0, 0, 0].map((n) => { return BigNumber.from(n) });
 
-const acc1BLSPubKey = [0xAA, 0xBB, 0xCC, 0xDD];
-const acc2BLSPubKey = [0xEE, 0xFF, 0x00, 0x11];
 const initialSupply = ethers.utils.parseUnits("1000000")
+const userStartAmount = initialSupply.div(2);
 
 describe('BLSWallet', async function () {
   beforeEach(async function () {
     [admin, account1, account2] = await ethers.getSigners();
 
+    // setup erc20 token and account balances
     const MockERC20 = await ethers.getContractFactory("MockERC20");
     this.baseToken = await MockERC20.deploy("AnyToken","TOK", initialSupply);
+    await this.baseToken.deployed();
+    await this.baseToken.transfer(account1.address, userStartAmount);
+    await this.baseToken.transfer(account2.address, userStartAmount);
 
-    const BLSWallet = await ethers.getContractFactory("BLSWallet");
+    // deploy bls wallet with token address
+    const BLSWallet = await ethers.getContractFactory("MockBLSWallet");
     this.blsWallet = await BLSWallet.deploy(this.baseToken.address);
+    await this.blsWallet.deployed();
+
+    // approve bls wallet with amount to transfer
+    await this.baseToken.connect(account1).approve(this.blsWallet.address, userStartAmount);
+    await this.baseToken.connect(account2).approve(this.blsWallet.address, userStartAmount);
   });
 
-  it('should ', async function () {
-    await this.blsWallet.connect(account1).deposit(acc1BLSPubKey, 0);
+  it('should deposit balance from token to bls wallet', async function () {
+    await this.blsWallet.connect(account1).deposit(acc1BLSPubKey, userStartAmount);
+    expect(await this.blsWallet.balanceOf(account1.address)).to.equal(userStartAmount);
   });
+
+  it('should set bls public key on deposit', async function () {
+    await this.blsWallet.connect(account1).deposit(acc1BLSPubKey, userStartAmount);
+    expect(await this.blsWallet.blsPubKeyOf(account1.address)).to.eql(acc1BLSPubKey);
+  });
+
+  it('should withdraw full balance from token to bls wallet', async function () {
+    await this.blsWallet.connect(account1).withdraw();
+    expect(await this.blsWallet.balanceOf(account1.address)).to.equal(0);
+  });
+
+  it('should reset bls public key on withdraw', async function () {
+    await this.blsWallet.connect(account1).withdraw();
+    expect(await this.blsWallet.blsPubKeyOf(account1.address)).to.eql(zeroBLSPubKey);
+  });
+
 });