From 16a8b020bc62d9b9b01919ef5c49df2df6078aee Mon Sep 17 00:00:00 2001
From: Mark Spanbroek <mark@spanbroek.net>
Date: Tue, 21 Jan 2025 13:37:01 +0100
Subject: [PATCH] vault: lock is deleted upon withdrawal

---
 contracts/Vault.sol |  2 ++
 test/Vault.tests.js | 10 ++++++++++
 2 files changed, 12 insertions(+)

diff --git a/contracts/Vault.sol b/contracts/Vault.sol
index 6afd24a..a67f630 100644
--- a/contracts/Vault.sol
+++ b/contracts/Vault.sol
@@ -67,7 +67,9 @@ contract Vault {
   }
 
   function withdraw(Context context, Recipient recipient) public {
+    Controller controller = Controller.wrap(msg.sender);
     require(!lock(context).expiry.isFuture(), Locked());
+    delete _locks[controller][context];
     uint256 amount = balance(context, recipient);
     _delete(context, recipient);
     _token.safeTransfer(Recipient.unwrap(recipient), amount);
diff --git a/test/Vault.tests.js b/test/Vault.tests.js
index 8a97dea..baf0913 100644
--- a/test/Vault.tests.js
+++ b/test/Vault.tests.js
@@ -359,5 +359,15 @@ describe("Vault", function () {
       const extending = vault.extend(context, maximum)
       await expect(extending).to.be.revertedWith("LockExpired")
     })
+
+    it("deletes lock when funds are withdrawn", async function () {
+      let start = await currentTime()
+      let expiry = start + 10
+      await vault.lockup(context, expiry, expiry)
+      await advanceTimeToForNextBlock(expiry)
+      await vault.withdraw(context, account.address)
+      expect((await vault.lock(context))[0]).to.equal(0)
+      expect((await vault.lock(context))[1]).to.equal(0)
+    })
   })
 })