Option to remove delegate

makerdao · Nov 24, 2023 · 5701a66 · 5701a66
1 parent d852143
commit 5701a66
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/src/LockstakeEngine.sol b/src/LockstakeEngine.sol
@@ -199,16 +199,18 @@ contract LockstakeEngine {
     }
 
     function delegate(address urn, address delegate_) external urnOwner(urn) {
-        require(delegateFactory.isDelegate(delegate_) == 1, "LockstakeEngine/not-valid-delegate");
+        require(delegate_ == address(0) || delegateFactory.isDelegate(delegate_) == 1, "LockstakeEngine/not-valid-delegate");
         address prevDelegate = urnDelegates[urn];
         require(prevDelegate != delegate_, "LockstakeEngine/same-delegate");
         (uint256 wad,) = vat.urns(ilk, urn);
         if (wad > 0) {
             if (prevDelegate != address(0)) {
                 DelegateLike(prevDelegate).free(wad);
             }
-            gov.approve(address(delegate_), wad);
-            DelegateLike(delegate_).lock(wad);
+            if (delegate_ != address(0)) {
+                gov.approve(address(delegate_), wad);
+                DelegateLike(delegate_).lock(wad);
+            }
         }
         urnDelegates[urn] = delegate_;
         emit Delegate(urn, delegate_);

diff --git a/test/LockstakeEngine.t.sol b/test/LockstakeEngine.t.sol
@@ -240,18 +240,29 @@ contract AllocatorVaultTest is DssTest {
     function testDelegate() public {
         address urn = engine.open();
         engine.delegate(urn, voterDelegate);
+        assertEq(engine.urnDelegates(urn), voterDelegate);
         vm.prank(address(888)); address voterDelegate2 = delFactory.create();
         gov.approve(address(engine), 100_000 * 10**18);
         engine.lock(urn, 100_000 * 10**18);
         assertEq(DelegateMock(voterDelegate).stake(address(engine)), 100_000 * 10**18);
         assertEq(DelegateMock(voterDelegate2).stake(address(engine)), 0);
         assertEq(gov.balanceOf(voterDelegate), 100_000 * 10**18);
         assertEq(gov.balanceOf(voterDelegate2), 0);
+        assertEq(gov.balanceOf(address(engine)), 0);
         engine.delegate(urn, voterDelegate2);
+        assertEq(engine.urnDelegates(urn), voterDelegate2);
         assertEq(DelegateMock(voterDelegate).stake(address(engine)), 0);
         assertEq(DelegateMock(voterDelegate2).stake(address(engine)), 100_000 * 10**18);
         assertEq(gov.balanceOf(voterDelegate), 0);
         assertEq(gov.balanceOf(voterDelegate2), 100_000 * 10**18);
+        assertEq(gov.balanceOf(address(engine)), 0);
+        engine.delegate(urn, address(0));
+        assertEq(engine.urnDelegates(urn), address(0));
+        assertEq(DelegateMock(voterDelegate).stake(address(engine)), 0);
+        assertEq(DelegateMock(voterDelegate2).stake(address(engine)), 0);
+        assertEq(gov.balanceOf(voterDelegate), 0);
+        assertEq(gov.balanceOf(voterDelegate2), 0);
+        assertEq(gov.balanceOf(address(engine)), 100_000 * 10**18);
     }
 
     function testDrawWipe() public {