re-enable some stake table tests

- Some of the tests do not apply anymore since the requirements changed.
  Notably anything relating to the queue logic for entering and exiting.
- Fix that we were not crediting delegations.
- Update undelegation test to test partial withdrawals.

contracts/src/ExampleToken.sol

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.0;
 
+// TODO MA: should we really use solmate, or just stick with openzeppelin?
 import { ERC20 } from "solmate/tokens/ERC20.sol";
 
 contract ExampleToken is ERC20 {

contracts/src/StakeTable.sol

@@ -180,8 +180,8 @@ contract StakeTable is AbstractStakeTable, Ownable, InitializedAt {
         ensureNonZeroSchnorrKey(schnorrVK);
         ensureNewKey(blsVK);
 
-        // Verify that the validator can sign for that blsVK.
-        // This prevents rogue public-key attacks.
+        // Verify that the validator can sign for that blsVK. This prevents rogue public-key
+        // attacks.
         bytes memory message = abi.encode(msg.sender);
         BLSSig.verifyBlsSig(message, blsSig, blsVK);
 
@@ -210,6 +210,7 @@ contract StakeTable is AbstractStakeTable, Ownable, InitializedAt {
         }
 
         validators[validator].delegatedAmount += amount;
+        delegations[validator][msg.sender] += amount;
 
         SafeTransferLib.safeTransferFrom(ERC20(tokenAddress), msg.sender, address(this), amount);
 
@@ -226,7 +227,7 @@ contract StakeTable is AbstractStakeTable, Ownable, InitializedAt {
 
         uint256 balance = delegations[validator][msg.sender];
         if (balance < amount) {
-            revert InsufficientBalance(amount);
+            revert InsufficientBalance(balance);
         }
 
         delegations[validator][msg.sender] -= amount;
@@ -291,14 +292,21 @@ contract StakeTable is AbstractStakeTable, Ownable, InitializedAt {
 
     /// @notice Update the consensus keys for a validator
     /// @dev This function is used to update the consensus keys for a validator
-    /// @dev This function can only be called by the validator itself when it's not in the exit
-    /// queue
+    /// @dev This function can only be called by the validator itself when it hasn't exited
+    ///      TODO: MA: is this a good idea? Why should key rotation be blocked for an exiting
+    ///      validator?
     /// @dev The validator will need to give up either its old BLS key and/or old Schnorr key
     /// @dev The validator will need to provide a BLS signature to prove that the account owns the
     /// new BLS key
     /// @param newBlsVK The new BLS verification key
     /// @param newSchnorrVK The new Schnorr verification key
     /// @param newBlsSig The BLS signature that the account owns the new BLS key
+    ///
+    /// TODO: MA: I think this function should be reworked. Is it fine to always force updating both
+    /// keys? If not we should probably rather have two functions for updating the keys. But this
+    /// would also mean two separate events, or storing the keys in the contract only for this
+    /// update function to remit the old keys, or throw errors if the keys are not changed. None of
+    /// that seems useful enough to warrant the extra complexity in the contract and GCL.
     function updateConsensusKeys(
         BN254.G2Point memory newBlsVK,
         EdOnBN254.EdOnBN254Point memory newSchnorrVK,
@@ -309,8 +317,8 @@ contract StakeTable is AbstractStakeTable, Ownable, InitializedAt {
         ensureNonZeroSchnorrKey(newSchnorrVK);
         ensureNewKey(newBlsVK);
 
-        // Verify that the validator can sign for that newBlsVK, otherwise it
-        // inner reverts with BLSSigVerificationFailed
+        // Verify that the validator can sign for that blsVK. This prevents rogue public-key
+        // attacks.
         bytes memory message = abi.encode(msg.sender);
         BLSSig.verifyBlsSig(message, newBlsSig, newBlsVK);