From dc7371bead16dd226e16ea9e9264eee8a76bad1c Mon Sep 17 00:00:00 2001
From: Yaru Wang <yaru@informal.systems>
Date: Mon, 2 Oct 2023 12:55:00 +0200
Subject: [PATCH] test: add unit test

---
 x/ccv/provider/keeper/gov_hook.go    |  3 +-
 x/ccv/provider/keeper/keeper.go      |  4 +-
 x/ccv/provider/keeper/keeper_test.go | 75 ++++++++++++++++++++++++++++
 3 files changed, 78 insertions(+), 4 deletions(-)

diff --git a/x/ccv/provider/keeper/gov_hook.go b/x/ccv/provider/keeper/gov_hook.go
index ca54584c13..61f504366f 100644
--- a/x/ccv/provider/keeper/gov_hook.go
+++ b/x/ccv/provider/keeper/gov_hook.go
@@ -82,9 +82,8 @@ func (gh GovHooks) AfterProposalVotingPeriodEnded(ctx sdk.Context, proposalID ui
 		}
 
 		if consAdditionProp.ProposalType() == types.ProposalTypeConsumerAddition {
-			gh.k.DeleteChainsInProposal(ctx, consAdditionProp.ChainId, proposalID)
+			gh.k.DeleteProposedConsumerChainInStore(ctx, proposalID)
 		}
-
 	}
 }
 
diff --git a/x/ccv/provider/keeper/keeper.go b/x/ccv/provider/keeper/keeper.go
index 92c3eefe4a..ba2b6f47fe 100644
--- a/x/ccv/provider/keeper/keeper.go
+++ b/x/ccv/provider/keeper/keeper.go
@@ -192,9 +192,9 @@ func (k Keeper) GetProposedConsumerChain(ctx sdk.Context, proposalID uint64) str
 	return string(store.Get(types.ProposedConsumerChainKey(proposalID)))
 }
 
-// DeleteChainsInProposal deletes the consumer chainID from store
+// DeleteProposedConsumerChainInStore deletes the consumer chainID from store
 // which is in gov consumerAddition proposal
-func (k Keeper) DeleteChainsInProposal(ctx sdk.Context, chainID string, proposalID uint64) {
+func (k Keeper) DeleteProposedConsumerChainInStore(ctx sdk.Context, proposalID uint64) {
 	store := ctx.KVStore(k.storeKey)
 	store.Delete(types.ProposedConsumerChainKey(proposalID))
 }
diff --git a/x/ccv/provider/keeper/keeper_test.go b/x/ccv/provider/keeper/keeper_test.go
index 12d693d389..73ef0af6ec 100644
--- a/x/ccv/provider/keeper/keeper_test.go
+++ b/x/ccv/provider/keeper/keeper_test.go
@@ -545,6 +545,7 @@ func TestSetProposedConsumerChains(t *testing.T) {
 		{chainID: "1", proposalID: 1},
 		{chainID: "some other ID", proposalID: 12},
 		{chainID: "some other other chain ID", proposalID: 123},
+		{chainID: "", proposalID: 1234},
 	}
 
 	for _, test := range tests {
@@ -553,3 +554,77 @@ func TestSetProposedConsumerChains(t *testing.T) {
 		require.Equal(t, cID, test.chainID)
 	}
 }
+
+func TestDeleteProposedConsumerChainInStore(t *testing.T) {
+	providerKeeper, ctx, ctrl, _ := testkeeper.GetProviderKeeperAndCtx(t, testkeeper.NewInMemKeeperParams(t))
+	defer ctrl.Finish()
+
+	tests := []struct {
+		chainID          string
+		proposalID       uint64
+		deleteProposalID uint64
+		ok               bool
+	}{
+		{chainID: "1", proposalID: 1, deleteProposalID: 1, ok: true},
+		{chainID: "", proposalID: 12, deleteProposalID: 12, ok: true},
+		{chainID: "1", proposalID: 0, deleteProposalID: 1, ok: false},
+	}
+	for _, test := range tests {
+		providerKeeper.SetProposedConsumerChain(ctx, test.chainID, test.proposalID)
+		providerKeeper.DeleteProposedConsumerChainInStore(ctx, test.deleteProposalID)
+		cID := providerKeeper.GetProposedConsumerChain(ctx, test.proposalID)
+		if test.ok {
+			require.Equal(t, cID, "")
+		} else {
+			require.Equal(t, cID, test.chainID)
+		}
+	}
+}
+
+func TestGetAllProposedConsumerChainIDs(t *testing.T) {
+	providerKeeper, ctx, ctrl, _ := testkeeper.GetProviderKeeperAndCtx(t, testkeeper.NewInMemKeeperParams(t))
+	defer ctrl.Finish()
+	tests := [][]types.ProposedChain{
+		[]types.ProposedChain{},
+		[]types.ProposedChain{
+			{
+				ChainID: "1",
+				ProposalID: 1,
+			},
+		},
+		[]types.ProposedChain{
+			{
+				ChainID: "1",
+				ProposalID: 1,
+			},
+			{
+				ChainID: "2",
+				ProposalID: 2,
+			},
+			{
+				ChainID: "",
+				ProposalID: 3,
+			},
+		},
+	}
+
+	for _, test := range tests {
+		for _, tc := range test {
+			providerKeeper.SetProposedConsumerChain(ctx, tc.ChainID, tc.ProposalID)
+		}
+
+		chains := providerKeeper.GetAllProposedConsumerChainIDs(ctx)
+
+		sort.Slice(chains, func(i, j int) bool {
+			return chains[i].ProposalID < chains[j].ProposalID
+		})
+		sort.Slice(test, func(i, j int) bool {
+			return test[i].ProposalID < test[j].ProposalID
+		})
+		require.Equal(t, chains, test)
+
+		for _, tc := range test {
+			providerKeeper.DeleteProposedConsumerChainInStore(ctx, tc.ProposalID)
+		}
+	}
+}