fix: query signal tally after successful try upgrade

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)

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)

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()