diff --git a/src/LockstakeEngine.sol b/src/LockstakeEngine.sol index e9edd6f7..cbfd6c2a 100644 --- a/src/LockstakeEngine.sol +++ b/src/LockstakeEngine.sol @@ -199,7 +199,7 @@ 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); @@ -207,8 +207,10 @@ contract LockstakeEngine { 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 index 2144edda..cb593c20 100644 --- a/test/LockstakeEngine.t.sol +++ b/test/LockstakeEngine.t.sol @@ -240,6 +240,7 @@ 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); @@ -247,11 +248,21 @@ contract AllocatorVaultTest is DssTest { 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 {