chore: add more state generator tests

pkg/fetchers/test_fetcher.go

@@ -8,6 +8,8 @@ import (
 type TestFetcher struct {
 	WithProposals      bool
 	WithProposalsError bool
+	WithVote           bool
+	WithVoteError      bool
 }
 
 func (f *TestFetcher) GetAllProposals(
@@ -34,6 +36,20 @@ func (f *TestFetcher) GetVote(
 	proposal, voter string,
 	prevHeight int64,
 ) (*types.Vote, int64, *types.QueryError) {
+	if f.WithVoteError {
+		return nil, 456, &types.QueryError{
+			QueryError: errors.New("error"),
+		}
+	}
+
+	if f.WithVote {
+		return &types.Vote{
+			ProposalID: "1",
+			Voter:      "me",
+			Options:    types.VoteOptions{},
+		}, 456, nil
+	}
+
 	return nil, 456, nil
 }
 

pkg/state/generator.go

@@ -120,33 +120,37 @@ func (g *Generator) ProcessProposalAndWallet(
 	oldVote, _, found := oldState.GetVoteAndProposal(chain.Name, proposal.ID, wallet.Address)
 	vote, voteHeight, err := fetcher.GetVote(proposal.ID, wallet.Address, oldVote.Height)
 
-	if found && oldVote.HasVoted() && vote == nil {
-		g.Logger.Trace().
-			Str("chain", chain.Name).
-			Str("proposal", proposal.ID).
-			Str("wallet", wallet.Address).
-			Msg("Wallet has voted and there's no vote in the new state - using old vote")
-
-		g.Mutex.Lock()
-		state.SetVote(
-			chain,
-			proposal,
-			wallet,
-			oldVote,
-		)
-		g.Mutex.Unlock()
-	}
-
 	proposalVote := ProposalVote{
 		Wallet: wallet,
 	}
 
 	if err != nil {
+		// 1. If error occurred - store the error, but preserve the older height and vote.
+		g.Logger.Trace().
+			Str("chain", chain.Name).
+			Str("proposal", proposal.ID).
+			Str("wallet", wallet.Address).
+			Int64("height", voteHeight).
+			Err(err).
+			Msg("Error fetching wallet vote - preserving the older height and vote")
+
 		proposalVote.Error = err
 		if found {
 			proposalVote.Height = oldVote.Height
+			proposalVote.Vote = oldVote.Vote
 		}
+	} else if found && oldVote.HasVoted() && vote == nil {
+		// 2. If there's no newer vote while there's an older vote - preserve the older vote
+		g.Logger.Trace().
+			Str("chain", chain.Name).
+			Str("proposal", proposal.ID).
+			Str("wallet", wallet.Address).
+			Msg("Wallet has voted and there's no vote in the new state - using old vote")
+
+		proposalVote.Vote = oldVote.Vote
+		proposalVote.Height = voteHeight
 	} else {
+		// 3. Wallet voted (or hadn't voted and hadn't voted before) - use the older vote.
 		proposalVote.Vote = vote
 		proposalVote.Height = voteHeight
 	}

pkg/state/generator_test.go

@@ -80,3 +80,219 @@ func TestReportGeneratorProcessProposalWithError(t *testing.T) {
 	assert.True(t, ok)
 	assert.Equal(t, "1", proposal.Proposal.ID)
 }
+
+func TestReportGeneratorProcessProposalWithoutError(t *testing.T) {
+	t.Parallel()
+
+	log := logger.GetNopLogger()
+	chain := &types.Chain{
+		Name:    "chain",
+		Type:    "cosmos",
+		Wallets: []*types.Wallet{{Address: "me"}},
+	}
+	chains := types.Chains{chain}
+	fetcher := &fetchers.TestFetcher{WithProposals: true, WithVote: true}
+
+	generator := Generator{
+		Logger: *log,
+		Chains: chains,
+		Fetchers: map[string]fetchers.Fetcher{
+			"chain": fetcher,
+		},
+	}
+
+	oldVotes := map[string]WalletVotes{
+		"1": {
+			Proposal: types.Proposal{ID: "1"},
+		},
+	}
+
+	oldState := NewState()
+	oldState.SetChainProposalsHeight(chain, 15)
+	oldState.SetChainVotes(chain, oldVotes)
+
+	newState := NewState()
+	generator.ProcessChain(chain, newState, oldState, fetcher)
+	assert.Len(t, newState.ChainInfos, 1)
+
+	newVotes, ok := newState.ChainInfos["chain"]
+	assert.True(t, ok)
+	assert.NotNil(t, newVotes)
+	assert.Nil(t, newVotes.ProposalsError)
+	assert.Equal(t, int64(123), newVotes.ProposalsHeight)
+
+	proposal, ok := newVotes.ProposalVotes["1"]
+	assert.True(t, ok)
+	assert.Equal(t, "1", proposal.Proposal.ID)
+}
+
+func TestReportGeneratorProcessVoteWithError(t *testing.T) {
+	t.Parallel()
+
+	log := logger.GetNopLogger()
+	wallet := &types.Wallet{Address: "me"}
+	chain := &types.Chain{
+		Name:    "chain",
+		Type:    "cosmos",
+		Wallets: []*types.Wallet{wallet},
+	}
+	chains := types.Chains{chain}
+	fetcher := &fetchers.TestFetcher{WithProposals: true, WithVoteError: true}
+
+	proposal := types.Proposal{ID: "1"}
+	generator := Generator{
+		Logger: *log,
+		Chains: chains,
+		Fetchers: map[string]fetchers.Fetcher{
+			"chain": fetcher,
+		},
+	}
+
+	oldVotes := map[string]WalletVotes{
+		"1": {
+			Proposal: proposal,
+			Votes: map[string]ProposalVote{
+				"me": {
+					Vote:   &types.Vote{Voter: "not_me"},
+					Height: 15,
+				},
+			},
+		},
+	}
+
+	oldState := NewState()
+	oldState.SetChainProposalsHeight(chain, 15)
+	oldState.SetChainVotes(chain, oldVotes)
+
+	newState := NewState()
+	generator.ProcessProposalAndWallet(chain, proposal, fetcher, wallet, newState, oldState)
+	assert.Len(t, newState.ChainInfos, 1)
+
+	newVotes, ok := newState.ChainInfos["chain"]
+	assert.True(t, ok)
+	assert.NotNil(t, newVotes)
+
+	newProposal, ok := newVotes.ProposalVotes["1"]
+	assert.True(t, ok)
+	assert.Equal(t, "1", newProposal.Proposal.ID)
+
+	newVote, ok := newProposal.Votes["me"]
+	assert.True(t, ok)
+	assert.Equal(t, int64(15), newVote.Height)
+	assert.NotNil(t, newVote.Error)
+	assert.Equal(t, "not_me", newVote.Vote.Voter)
+}
+
+func TestReportGeneratorProcessVoteWithDisappearedVote(t *testing.T) {
+	t.Parallel()
+
+	log := logger.GetNopLogger()
+	wallet := &types.Wallet{Address: "me"}
+	chain := &types.Chain{
+		Name:    "chain",
+		Type:    "cosmos",
+		Wallets: []*types.Wallet{wallet},
+	}
+	chains := types.Chains{chain}
+	fetcher := &fetchers.TestFetcher{WithProposals: true}
+
+	proposal := types.Proposal{ID: "1"}
+	generator := Generator{
+		Logger: *log,
+		Chains: chains,
+		Fetchers: map[string]fetchers.Fetcher{
+			"chain": fetcher,
+		},
+	}
+
+	oldVotes := map[string]WalletVotes{
+		"1": {
+			Proposal: proposal,
+			Votes: map[string]ProposalVote{
+				"me": {
+					Vote:   &types.Vote{Voter: "not_me"},
+					Height: 15,
+				},
+			},
+		},
+	}
+
+	oldState := NewState()
+	oldState.SetChainProposalsHeight(chain, 15)
+	oldState.SetChainVotes(chain, oldVotes)
+
+	newState := NewState()
+	generator.ProcessProposalAndWallet(chain, proposal, fetcher, wallet, newState, oldState)
+	assert.Len(t, newState.ChainInfos, 1)
+
+	newVotes, ok := newState.ChainInfos["chain"]
+	assert.True(t, ok)
+	assert.NotNil(t, newVotes)
+
+	newProposal, ok := newVotes.ProposalVotes["1"]
+	assert.True(t, ok)
+	assert.Equal(t, "1", newProposal.Proposal.ID)
+
+	newVote, ok := newProposal.Votes["me"]
+	assert.True(t, ok)
+	assert.Equal(t, int64(456), newVote.Height)
+	assert.Nil(t, newVote.Error)
+	assert.Equal(t, "not_me", newVote.Vote.Voter)
+}
+
+func TestReportGeneratorProcessVoteWithOkVote(t *testing.T) {
+	t.Parallel()
+
+	log := logger.GetNopLogger()
+	wallet := &types.Wallet{Address: "me"}
+	chain := &types.Chain{
+		Name:    "chain",
+		Type:    "cosmos",
+		Wallets: []*types.Wallet{wallet},
+	}
+	chains := types.Chains{chain}
+	fetcher := &fetchers.TestFetcher{WithProposals: true, WithVote: true}
+
+	proposal := types.Proposal{ID: "1"}
+	generator := Generator{
+		Logger: *log,
+		Chains: chains,
+		Fetchers: map[string]fetchers.Fetcher{
+			"chain": fetcher,
+		},
+	}
+
+	oldVotes := map[string]WalletVotes{
+		"1": {
+			Proposal: proposal,
+			Votes: map[string]ProposalVote{
+				"me": {
+					Vote:   &types.Vote{Voter: "not_me"},
+					Height: 15,
+				},
+			},
+		},
+	}
+
+	oldState := NewState()
+	oldState.SetChainProposalsHeight(chain, 15)
+	oldState.SetChainVotes(chain, oldVotes)
+
+	newState := NewState()
+	generator.ProcessProposalAndWallet(chain, proposal, fetcher, wallet, newState, oldState)
+	assert.Len(t, newState.ChainInfos, 1)
+
+	newVotes, ok := newState.ChainInfos["chain"]
+	assert.True(t, ok)
+	assert.NotNil(t, newVotes)
+
+	newProposal, ok := newVotes.ProposalVotes["1"]
+	assert.True(t, ok)
+	assert.Equal(t, "1", newProposal.Proposal.ID)
+
+	newVote, ok := newProposal.Votes["me"]
+	assert.True(t, ok)
+	assert.Equal(t, int64(456), newVote.Height)
+	assert.Nil(t, newVote.Error)
+	assert.Equal(t, "me", newVote.Vote.Voter)
+}