From e8b841080dd89cd4bb6596363502208bf8116d54 Mon Sep 17 00:00:00 2001
From: Christopher Goes <cwgoes@pluranimity.org>
Date: Thu, 5 Jul 2018 20:21:56 +0200
Subject: [PATCH] Update .Slash() and testcase

---
 x/stake/keeper/slash.go          |  6 +++++-
 x/stake/keeper/validator_test.go | 25 +++++++------------------
 2 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/x/stake/keeper/slash.go b/x/stake/keeper/slash.go
index b12360fc3f..2f620d9d47 100644
--- a/x/stake/keeper/slash.go
+++ b/x/stake/keeper/slash.go
@@ -92,7 +92,11 @@ func (k Keeper) Slash(ctx sdk.Context, pubkey crypto.PubKey, infractionHeight in
 	// update the pool
 	k.SetPool(ctx, pool)
 	// update the validator, possibly kicking it out
-	k.UpdateValidator(ctx, validator)
+	validator = k.UpdateValidator(ctx, validator)
+	// remove validator if it has been reduced to zero shares
+	if validator.PoolShares.Amount.IsZero() {
+		k.RemoveValidator(ctx, validator.Owner)
+	}
 
 	// Log that a slash occurred!
 	logger.Info(fmt.Sprintf("Validator %s slashed by slashFactor %v, removed %v shares and burned %d tokens", pubkey.Address(), slashFactor, sharesToRemove, burned))
diff --git a/x/stake/keeper/validator_test.go b/x/stake/keeper/validator_test.go
index c7607dd7e9..c037948da0 100644
--- a/x/stake/keeper/validator_test.go
+++ b/x/stake/keeper/validator_test.go
@@ -99,36 +99,25 @@ func TestZeroPowerUnbonded(t *testing.T) {
 
 	// create a random pool
 	pool.LooseTokens = 10000
-	pool.BondedTokens = 1234
-	pool.BondedShares = sdk.NewRat(124)
-	pool.UnbondingTokens = 13934
-	pool.UnbondingShares = sdk.NewRat(145)
-	pool.UnbondedTokens = 154
-	pool.UnbondedShares = sdk.NewRat(1333)
 	keeper.SetPool(ctx, pool)
 
 	// add a validator
 	validator := types.NewValidator(addrVals[0], PKs[0], types.Description{})
-	validator, pool, delSharesCreated := validator.AddTokensFromDel(pool, 100)
+	validator, pool, _ = validator.AddTokensFromDel(pool, 100)
 	require.Equal(t, sdk.Unbonded, validator.Status())
 	require.Equal(t, int64(100), validator.PoolShares.Tokens(pool).RoundInt64())
 	keeper.SetPool(ctx, pool)
+	keeper.SetValidatorByPubKeyIndex(ctx, validator)
 	keeper.UpdateValidator(ctx, validator)
 	validator, found := keeper.GetValidator(ctx, addrVals[0])
 	require.True(t, found)
 	require.Equal(t, int64(100), validator.PoolShares.Tokens(pool).RoundInt64(), "\nvalidator %v\npool %v", validator, pool)
 
-	pool = keeper.GetPool(ctx)
-
-	validator, pool, burned := validator.RemoveDelShares(pool, delSharesCreated)
-	require.Equal(t, int64(100), burned)
-	keeper.SetPool(ctx, pool)              // update the pool
-	keeper.UpdateValidator(ctx, validator) // update the validator, kicking it out
-	pool = keeper.GetPool(ctx)
-	validator, found = keeper.GetValidator(ctx, addrVals[0])
-	require.True(t, found)
-	require.Equal(t, sdk.ZeroRat(), validator.GetPower())
-	require.Equal(t, sdk.Unbonded, validator.GetStatus())
+	// slash the validator by 100%
+	keeper.Slash(ctx, PKs[0], 0, 100, sdk.OneRat())
+	// validator should have been deleted
+	_, found = keeper.GetValidator(ctx, addrVals[0])
+	require.False(t, found)
 }
 
 // This function tests UpdateValidator, GetValidator, GetValidatorsBonded, RemoveValidator