From 75e499dbb8836e8d48dcf95e168fb6019e906a6c Mon Sep 17 00:00:00 2001
From: Yarom Swisa <yarom@lavanet.xyz git config --global user.name Yarom>
Date: Mon, 23 Sep 2024 16:15:59 +0300
Subject: [PATCH] adding tests and fix

---
 x/pairing/keeper/grpc_query_provider.go  | 11 ++++---
 x/pairing/keeper/single_provider_test.go | 39 ++++++++++++++++++++++++
 x/pairing/keeper/unstaking.go            |  1 +
 3 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/x/pairing/keeper/grpc_query_provider.go b/x/pairing/keeper/grpc_query_provider.go
index d2049d0f63..8296712dbf 100644
--- a/x/pairing/keeper/grpc_query_provider.go
+++ b/x/pairing/keeper/grpc_query_provider.go
@@ -17,15 +17,16 @@ func (k Keeper) Provider(goCtx context.Context, req *types.QueryProviderRequest)
 
 	ctx := sdk.UnwrapSDKContext(goCtx)
 
+	metadata, err := k.epochStorageKeeper.GetMetadata(ctx, req.Address)
+	if err != nil {
+		return &types.QueryProviderResponse{}, nil
+	}
+
 	chains := []string{req.ChainID}
 	if req.ChainID == "" {
-		chains = k.specKeeper.GetAllChainIDs(ctx)
+		chains = metadata.Chains
 	}
 
-	metadata, err := k.epochStorageKeeper.GetMetadata(ctx, req.Address)
-	if err != nil {
-		return nil, err
-	}
 	stakeEntries := []epochstoragetypes.StakeEntry{}
 	for _, chain := range chains {
 		stakeEntry, found := k.epochStorageKeeper.GetStakeEntryCurrent(ctx, chain, req.Address)
diff --git a/x/pairing/keeper/single_provider_test.go b/x/pairing/keeper/single_provider_test.go
index 41732eb4e2..34546b30dd 100644
--- a/x/pairing/keeper/single_provider_test.go
+++ b/x/pairing/keeper/single_provider_test.go
@@ -162,3 +162,42 @@ func TestDelegations(t *testing.T) {
 		require.Equal(t, int64(0), res.StakeEntries[0].DelegateTotal.Amount.Int64())
 	}
 }
+
+func TestUnstakeWithOperator(t *testing.T) {
+	ts := newTester(t)
+	SetupForSingleProviderTests(ts, 1, 5, 0)
+
+	// delegator, _ := ts.AddAccount("del", 1, 1000000000)
+	provider, _ := ts.GetAccount(common.PROVIDER, 0)
+
+	// unstake one chain
+	_, err := ts.TxPairingUnstakeProvider(provider.Addr.String(), SpecName(0))
+	require.NoError(t, err)
+
+	// check the entries got the stake of spec0
+	for i := 1; i < 5; i++ {
+		res, err := ts.QueryPairingProvider(provider.Addr.String(), SpecName(i))
+		require.NoError(t, err)
+		require.Equal(t, testStake*5/4, res.StakeEntries[0].Stake.Amount.Int64())
+	}
+
+	// unstake everything
+	for i := 1; i < 5; i++ {
+		_, err := ts.TxPairingUnstakeProvider(provider.Addr.String(), SpecName(i))
+		require.NoError(t, err)
+	}
+
+	res, err := ts.QueryPairingProvider(provider.Addr.String(), "")
+	require.NoError(t, err)
+	require.Len(t, res.StakeEntries, 0)
+
+	// stake again
+	d := common.MockDescription()
+	err = ts.StakeProviderExtra(provider.GetVaultAddr(), provider.Addr.String(), ts.Spec(SpecName(0)), 1, nil, 0, d.Moniker, d.Identity, d.Website, d.SecurityContact, d.Details)
+	require.NoError(ts.T, err)
+
+	// check we got all the self delegations
+	res, err = ts.QueryPairingProvider(provider.Addr.String(), SpecName(0))
+	require.NoError(t, err)
+	require.Equal(t, testStake*5+1, res.StakeEntries[0].Stake.Amount.Int64())
+}
diff --git a/x/pairing/keeper/unstaking.go b/x/pairing/keeper/unstaking.go
index 3b77fbb1da..d44f860546 100644
--- a/x/pairing/keeper/unstaking.go
+++ b/x/pairing/keeper/unstaking.go
@@ -77,6 +77,7 @@ func (k Keeper) UnstakeEntry(ctx sdk.Context, validator, chainID, creator, unsta
 				entry.Stake = entry.Stake.AddAmount(part)
 				total = total.Sub(part)
 				count--
+				k.epochStorageKeeper.SetStakeEntryCurrent(ctx, entry)
 			}
 		}
 	}