Prevent withdrawal of locked collateral

contracts/Collateral.sol

@@ -2,8 +2,9 @@
 pragma solidity ^0.8.0;
 
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "./AccountLocks.sol";
 
-contract Collateral {
+contract Collateral is AccountLocks {
   IERC20 private immutable token;
   Totals private totals;
   mapping(address => uint256) private balances;
@@ -38,6 +39,7 @@ contract Collateral {
   }
 
   function withdraw() public invariant {
+    _unlockAccount();
     uint256 amount = balanceOf(msg.sender);
     totals.withdrawn += amount;
     subtract(msg.sender, amount);

contracts/TestCollateral.sol

@@ -11,4 +11,16 @@ contract TestCollateral is Collateral {
   function slash(address account, uint256 percentage) public {
     _slash(account, percentage);
   }
+
+  function createLock(bytes32 id, uint256 expiry) public {
+    _createLock(id, expiry);
+  }
+
+  function lock(address account, bytes32 id) public {
+    _lock(account, id);
+  }
+
+  function unlock(bytes32 id) public {
+    _unlock(id);
+  }
 }

test/Collateral.test.js

@@ -1,4 +1,5 @@
 const { expect } = require("chai")
+const { exampleLock } = require("./examples")
 
 describe("Collateral", function () {
   let collateral, token
@@ -89,4 +90,25 @@ describe("Collateral", function () {
       expect(await collateral.balanceOf(account1.address)).to.equal(950)
     })
   })
+
+  describe("locking", function () {
+    let lock
+
+    beforeEach(async function () {
+      await token.approve(collateral.address, 42)
+      await collateral.deposit(42)
+      lock = exampleLock()
+      await collateral.createLock(lock.id, lock.expiry)
+      await collateral.lock(account0.address, lock.id)
+    })
+
+    it("withdrawal fails when account is locked", async function () {
+      await expect(collateral.withdraw()).to.be.revertedWith("Account locked")
+    })
+
+    it("withdrawal succeeds when account is unlocked", async function () {
+      await collateral.unlock(lock.id)
+      await expect(collateral.withdraw()).not.to.be.reverted
+    })
+  })
 })