From 0bfd074807b6e9a4c882f386432f4a4ea85f3478 Mon Sep 17 00:00:00 2001
From: Rootul P <rootulp@gmail.com>
Date: Mon, 25 Nov 2024 08:51:24 -0500
Subject: [PATCH] fix: `query signal tally` after successful try upgrade
 (#4045)

Closes https://github.com/celestiaorg/celestia-app/issues/4007

## Testing

Using the updated single-node.sh script (see other PR), I can query the
version tally even after a successful try upgrade.

```
$ ./scripts/upgrade-to-v3.sh

$ celestia-appd query signal tally 3
threshold_power: "4167"
total_voting_power: "5000"
voting_power: "5000"
```
---
 x/signal/integration_test.go | 10 ++++++++++
 x/signal/keeper.go           |  7 +++++++
 x/signal/keeper_test.go      | 36 ++++++++++++++++++++++++++++++++++++
 3 files changed, 53 insertions(+)

diff --git a/x/signal/integration_test.go b/x/signal/integration_test.go
index 829a2d0be9..e1f79db39b 100644
--- a/x/signal/integration_test.go
+++ b/x/signal/integration_test.go
@@ -60,6 +60,16 @@ func TestUpgradeIntegration(t *testing.T) {
 	_, err = app.SignalKeeper.TryUpgrade(ctx, nil)
 	require.NoError(t, err)
 
+	// Verify that if a user queries the version tally, it still works after a
+	// successful try upgrade.
+	res, err = app.SignalKeeper.VersionTally(ctx, &types.QueryVersionTallyRequest{
+		Version: 3,
+	})
+	require.NoError(t, err)
+	require.EqualValues(t, 1, res.VotingPower)
+	require.EqualValues(t, 1, res.ThresholdPower)
+	require.EqualValues(t, 1, res.TotalVotingPower)
+
 	// Verify that if a subsequent call to TryUpgrade is made, it returns an
 	// error because an upgrade is already pending.
 	_, err = app.SignalKeeper.TryUpgrade(ctx, nil)
diff --git a/x/signal/keeper.go b/x/signal/keeper.go
index 7e85e604d5..a53a5a8c2a 100644
--- a/x/signal/keeper.go
+++ b/x/signal/keeper.go
@@ -1,6 +1,7 @@
 package signal
 
 import (
+	"bytes"
 	"context"
 	"encoding/binary"
 
@@ -122,6 +123,9 @@ func (k Keeper) VersionTally(ctx context.Context, req *types.QueryVersionTallyRe
 	iterator := store.Iterator(types.FirstSignalKey, nil)
 	defer iterator.Close()
 	for ; iterator.Valid(); iterator.Next() {
+		if bytes.Equal(iterator.Key(), types.UpgradeKey) {
+			continue
+		}
 		valAddress := sdk.ValAddress(iterator.Key())
 		power := k.stakingKeeper.GetLastValidatorPower(sdkCtx, valAddress)
 		version := VersionFromBytes(iterator.Value())
@@ -158,6 +162,9 @@ func (k Keeper) TallyVotingPower(ctx sdk.Context, threshold int64) (bool, uint64
 	iterator := store.Iterator(types.FirstSignalKey, nil)
 	defer iterator.Close()
 	for ; iterator.Valid(); iterator.Next() {
+		if bytes.Equal(iterator.Key(), types.UpgradeKey) {
+			continue
+		}
 		valAddress := sdk.ValAddress(iterator.Key())
 		// check that the validator is still part of the bonded set
 		val, found := k.stakingKeeper.GetValidator(ctx, valAddress)
diff --git a/x/signal/keeper_test.go b/x/signal/keeper_test.go
index 6c688a3328..676c93f838 100644
--- a/x/signal/keeper_test.go
+++ b/x/signal/keeper_test.go
@@ -430,6 +430,42 @@ func TestGetUpgrade(t *testing.T) {
 	})
 }
 
+func TestTallyAfterTryUpgrade(t *testing.T) {
+	upgradeKeeper, ctx, _ := setup(t)
+	goCtx := sdk.WrapSDKContext(ctx)
+
+	_, err := upgradeKeeper.SignalVersion(goCtx, &types.MsgSignalVersion{
+		ValidatorAddress: testutil.ValAddrs[0].String(),
+		Version:          3,
+	})
+	require.NoError(t, err)
+
+	_, err = upgradeKeeper.SignalVersion(goCtx, &types.MsgSignalVersion{
+		ValidatorAddress: testutil.ValAddrs[1].String(),
+		Version:          3,
+	})
+	require.NoError(t, err)
+
+	_, err = upgradeKeeper.SignalVersion(goCtx, &types.MsgSignalVersion{
+		ValidatorAddress: testutil.ValAddrs[2].String(),
+		Version:          3,
+	})
+	require.NoError(t, err)
+
+	_, err = upgradeKeeper.TryUpgrade(goCtx, &types.MsgTryUpgrade{})
+	require.NoError(t, err)
+
+	// Previously there was a bug where querying for the version tally after a
+	// successful try upgrade would result in a panic. See
+	// https://github.com/celestiaorg/celestia-app/issues/4007
+	res, err := upgradeKeeper.VersionTally(goCtx, &types.QueryVersionTallyRequest{
+		Version: 2,
+	})
+	require.NoError(t, err)
+	require.EqualValues(t, 100, res.ThresholdPower)
+	require.EqualValues(t, 120, res.TotalVotingPower)
+}
+
 func setup(t *testing.T) (signal.Keeper, sdk.Context, *mockStakingKeeper) {
 	signalStore := sdk.NewKVStoreKey(types.StoreKey)
 	db := tmdb.NewMemDB()