Fix: Staking/RewardPool contracts

packages/core/contracts/RewardPool.sol

@@ -106,6 +106,8 @@ contract RewardPool is IRewardPool, OwnableUpgradeable, UUPSUpgradeable {
     function distributeReward(address _escrowAddress) external override {
         Reward[] memory rewardsForEscrow = rewards[_escrowAddress];
 
+        require(rewardsForEscrow.length > 0, 'No rewards for escrow');
+
         // Delete rewards for allocation
         delete rewards[_escrowAddress];
 

packages/core/contracts/Staking.sol

@@ -282,7 +282,8 @@ contract Staking is IStaking, OwnableUpgradeable, UUPSUpgradeable {
 
         if (
             allocation.closedAt == 0 &&
-            escrowStatus == IEscrow.EscrowStatuses.Complete
+            (escrowStatus == IEscrow.EscrowStatuses.Complete ||
+                escrowStatus == IEscrow.EscrowStatuses.Cancelled)
         ) {
             return AllocationState.Completed;
         }
@@ -371,7 +372,6 @@ contract Staking is IStaking, OwnableUpgradeable, UUPSUpgradeable {
     function unstake(uint256 _tokens) external override onlyStaker(msg.sender) {
         Stakes.Staker storage staker = stakes[msg.sender];
 
-        require(staker.tokensStaked > 0, 'Must be a positive number');
         require(_tokens > 0, 'Must be a positive number');
         require(
             staker.tokensAvailable() >= _tokens,
@@ -439,7 +439,7 @@ contract Staking is IStaking, OwnableUpgradeable, UUPSUpgradeable {
 
         Allocation storage allocation = allocations[_escrowAddress];
 
-        require(allocation.tokens > 0, 'Must be a positive number');
+        require(_tokens > 0, 'Must be a positive number');
 
         require(
             _tokens <= allocation.tokens,
@@ -519,6 +519,8 @@ contract Staking is IStaking, OwnableUpgradeable, UUPSUpgradeable {
     function closeAllocation(
         address _escrowAddress
     ) external override onlyStaker(msg.sender) {
+        require(_escrowAddress != address(0), 'Must be a valid address');
+
         _closeAllocation(_escrowAddress);
     }
 
@@ -533,7 +535,7 @@ contract Staking is IStaking, OwnableUpgradeable, UUPSUpgradeable {
             'Allocation has no completed state'
         );
 
-        Allocation memory allocation = allocations[_escrowAddress];
+        Allocation storage allocation = allocations[_escrowAddress];
 
         allocation.closedAt = block.number;
         uint256 diffInBlocks = Math.diffOrZero(
@@ -542,11 +544,14 @@ contract Staking is IStaking, OwnableUpgradeable, UUPSUpgradeable {
         );
         require(diffInBlocks > 0, 'Allocation cannot be closed so early');
 
-        stakes[allocation.staker].unallocate(allocation.tokens);
+        uint256 _tokens = allocation.tokens;
+
+        stakes[allocation.staker].unallocate(_tokens);
+        allocation.tokens = 0;
 
         emit AllocationClosed(
             allocation.staker,
-            allocation.tokens,
+            _tokens,
             _escrowAddress,
             allocation.closedAt
         );

packages/core/test/RewardPool.ts

@@ -229,6 +229,12 @@ describe('RewardPool', function () {
       await staking.connect(operator).allocate(escrowAddress, allocatedTokens);
     });
 
+    it('Should revert if there is no reward', async () => {
+      await expect(
+        rewardPool.distributeReward(escrowAddress)
+      ).to.be.revertedWith('No rewards for escrow');
+    });
+
     it('Should distribute the reward.', async () => {
       const vSlashAmount = 2;
       const v2SlashAmount = 3;

packages/core/test/Staking.ts

@@ -3,12 +3,16 @@ import { expect } from 'chai';
 import { ethers, upgrades } from 'hardhat';
 import { Signer } from 'ethers';
 import {
+  Escrow,
   EscrowFactory,
   HMToken,
   Staking,
   RewardPool,
 } from '../typechain-types';
 
+const MOCK_URL = 'http://google.com/fake';
+const MOCK_HASH = 'kGKmnj9BRf';
+
 const mineNBlocks = async (n: number) => {
   await Promise.all(
     Array(n)
@@ -187,6 +191,14 @@ describe('Staking', function () {
         ).to.be.revertedWith('Must be a positive number');
       });
 
+      it('Should revert with the right error if total stake is below the unstake amount', async function () {
+        const amount = 15;
+
+        await expect(
+          staking.connect(operator).unstake(amount)
+        ).to.be.revertedWith('Insufficient amount to unstake');
+      });
+
       it('Should revert with the right error if total stake is below the minimum threshold', async function () {
         const amount = 9;
 
@@ -708,7 +720,19 @@ describe('Staking', function () {
             )
         ).to.be.revertedWith('Slash tokens exceed allocated ones');
       });
-      // TODO: Add additional tests
+
+      it('Should revert if slash amount is 0', async function () {
+        await expect(
+          staking
+            .connect(owner)
+            .slash(
+              await validator.getAddress(),
+              await operator.getAddress(),
+              escrowAddress,
+              0
+            )
+        ).to.be.revertedWith('Must be a positive number');
+      });
     });
 
     describe('Events', function () {
@@ -810,4 +834,118 @@ describe('Staking', function () {
       );
     });
   });
+
+  describe('closeAllocation', function () {
+    let escrowAddress: string;
+    let escrow: Escrow;
+
+    this.beforeEach(async () => {
+      const amount = 10;
+      const allocationAmount = 5;
+
+      await staking.connect(operator).stake(amount);
+
+      const result = await (
+        await escrowFactory
+          .connect(operator)
+          .createEscrow(token.address, [
+            await validator.getAddress(),
+            await reputationOracle.getAddress(),
+            await recordingOracle.getAddress(),
+          ])
+      ).wait();
+      const event = result.events?.find(({ topics }) =>
+        topics.includes(ethers.utils.id('Launched(address,address)'))
+      )?.args;
+
+      expect(event?.token).to.equal(token.address, 'token address is correct');
+      expect(event?.escrow).to.not.be.null;
+
+      escrowAddress = event?.escrow;
+
+      // Fund escrow
+      await token.connect(owner).transfer(escrowAddress, 100);
+
+      const EscrowFactory = await ethers.getContractFactory('Escrow');
+      escrow = await EscrowFactory.attach(escrowAddress);
+
+      // Setup escrow
+      await escrow
+        .connect(operator)
+        .setup(
+          await reputationOracle.getAddress(),
+          await recordingOracle.getAddress(),
+          10,
+          10,
+          MOCK_URL,
+          MOCK_HASH
+        );
+
+      await staking
+        .connect(operator)
+        .allocate(escrowAddress.toString(), allocationAmount);
+    });
+
+    describe('Validations', function () {
+      it('Should revert with the right error if not a valid address', async function () {
+        await expect(
+          staking
+            .connect(operator)
+            .closeAllocation(ethers.constants.AddressZero)
+        ).to.be.revertedWith('Must be a valid address');
+      });
+
+      it('Should revert with the right error if escrow is not completed nor cancelled', async function () {
+        await expect(
+          staking.connect(operator).closeAllocation(escrowAddress.toString())
+        ).to.be.revertedWith('Allocation has no completed state');
+      });
+    });
+
+    describe('Close allocation on completed/cancelled escrows', function () {
+      it('Should close allocation on completed escrows', async function () {
+        // Bulk payout & Complete Escrow
+        await escrow
+          .connect(operator)
+          .bulkPayOut(
+            [await operator2.getAddress()],
+            [100],
+            MOCK_URL,
+            MOCK_HASH,
+            '000'
+          );
+
+        await escrow.connect(operator).complete();
+
+        // Close allocation
+        await staking
+          .connect(operator)
+          .closeAllocation(escrowAddress.toString());
+
+        const allocation = await staking
+          .connect(operator)
+          .getAllocation(escrowAddress.toString());
+
+        expect(allocation.closedAt).not.to.be.null;
+        expect(allocation.tokens.toString()).to.equal('0');
+      });
+
+      it('Should close allocation on cancelled escrows', async function () {
+        // Close escrow
+        await escrow.connect(operator).cancel();
+
+        // Close allocation
+        await staking
+          .connect(operator)
+          .closeAllocation(escrowAddress.toString());
+
+        const allocation = await staking
+          .connect(operator)
+          .getAllocation(escrowAddress.toString());
+
+        expect(allocation.closedAt).not.to.be.null;
+        expect(allocation.tokens.toString()).to.equal('0');
+      });
+    });
+  });
 });