go/staking: Fix incorrect non-existence check

oasisprotocol · Nov 3, 2021 · f7c17e1 · f7c17e1
1 parent aa0e4ef
commit f7c17e1
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/go/staking/api/api.go b/go/staking/api/api.go
@@ -835,7 +835,7 @@ func (sa *StakeAccumulator) AddClaimUnchecked(claim StakeClaim, thresholds []Sta
 //
 // It is an error if the stake claim does not exist.
 func (sa *StakeAccumulator) RemoveClaim(claim StakeClaim) error {
-	if sa.Claims == nil || sa.Claims[claim] == nil {
+	if _, exists := sa.Claims[claim]; !exists {
 		return fmt.Errorf("staking: claim does not exist: %s", claim)
 	}
 

diff --git a/go/staking/api/api_test.go b/go/staking/api/api_test.go
@@ -206,6 +206,17 @@ func TestStakeAccumulator(t *testing.T) {
 	err = acct.CheckStakeClaims(thresholds)
 	require.NoError(err, "escrow account should check out")
 
+	// Add claim with empty threshold list.
+	err = acct.AddStakeClaim(thresholds, StakeClaim("claimEmptyList"), []StakeThreshold{})
+	require.NoError(err, "adding an empty list stake claim should work")
+	err = acct.RemoveStakeClaim(StakeClaim("claimEmptyList"))
+	require.NoError(err, "removing an empty claim should work")
+
+	err = acct.AddStakeClaim(thresholds, StakeClaim("claimNilList"), nil)
+	require.NoError(err, "adding an nil list stake claim should work")
+	err = acct.RemoveStakeClaim(StakeClaim("claimNilList"))
+	require.NoError(err, "removing an nil claim should work")
+
 	// Reduce stake.
 	acct.Active.Balance = *quantity.NewFromUint64(5_000)
 	err = acct.CheckStakeClaims(thresholds)