chore: add state generator test (#78)

* chore: add state generator test

* chore: add more state generator tests

pkg/fetchers/test_fetcher.go

@@ -0,0 +1,62 @@
+package fetchers
+
+import (
+	"errors"
+	"main/pkg/types"
+)
+
+type TestFetcher struct {
+	WithProposals      bool
+	WithProposalsError bool
+	WithVote           bool
+	WithVoteError      bool
+}
+
+func (f *TestFetcher) GetAllProposals(
+	prevHeight int64,
+) ([]types.Proposal, int64, *types.QueryError) {
+	if f.WithProposalsError {
+		return []types.Proposal{}, 123, &types.QueryError{
+			QueryError: errors.New("error"),
+		}
+	}
+
+	if f.WithProposals {
+		return []types.Proposal{
+			{
+				ID: "1",
+			},
+		}, 123, nil
+	}
+
+	return []types.Proposal{}, 123, nil
+}
+
+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
+}
+
+func (f *TestFetcher) GetTallies() (types.ChainTallyInfos, error) {
+	return types.ChainTallyInfos{}, nil
+}
+
+func (f *TestFetcher) GetChainParams() (*types.ChainWithVotingParams, []error) {
+	return nil, []error{}
+}

pkg/fs/test_fs.go

@@ -0,0 +1,30 @@
+package fs
+
+import (
+	"os"
+)
+
+type TestFS struct{}
+
+type TestFile struct {
+}
+
+func (f *TestFile) Write(p []byte) (int, error) {
+	return 0, nil
+}
+
+func (f *TestFile) Close() error {
+	return nil
+}
+
+func (fs *TestFS) ReadFile(name string) ([]byte, error) {
+	return []byte{}, nil
+}
+
+func (fs *TestFS) WriteFile(name string, data []byte, perms os.FileMode) error {
+	return nil
+}
+
+func (fs *TestFS) Create(path string) (File, error) {
+	return &TestFile{}, nil // go
+}

pkg/logger/logger.go

@@ -12,6 +12,11 @@ func GetDefaultLogger() *zerolog.Logger {
 	return &log
 }
 
+func GetNopLogger() *zerolog.Logger {
+	log := zerolog.Nop()
+	return &log
+}
+
 func GetLogger(config types.LogConfig) *zerolog.Logger {
 	log := zerolog.New(zerolog.ConsoleWriter{Out: os.Stdout}).With().Timestamp().Logger()
 

pkg/report/generator_test.go

@@ -4,7 +4,6 @@ import (
 	"errors"
 	"main/pkg/events"
 	"main/pkg/fs"
-	"os"
 	"testing"
 
 	"main/pkg/logger"
@@ -14,24 +13,10 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-type TestFS struct{}
-
-func (fs *TestFS) ReadFile(name string) ([]byte, error) {
-	return []byte{}, nil
-}
-
-func (fs *TestFS) WriteFile(name string, data []byte, perms os.FileMode) error {
-	return nil
-}
-
-func (fs *TestFS) Create(path string) (fs.File, error) {
-	return nil, nil
-}
-
 func TestReportGeneratorWithProposalError(t *testing.T) {
 	t.Parallel()
 
-	stateManager := state.NewStateManager("./state.json", &TestFS{}, logger.GetDefaultLogger())
+	stateManager := state.NewStateManager("./state.json", &fs.TestFS{}, logger.GetNopLogger())
 
 	oldState := state.NewState()
 	newState := state.State{
@@ -44,7 +29,7 @@ func TestReportGeneratorWithProposalError(t *testing.T) {
 		},
 	}
 
-	generator := NewReportGenerator(stateManager, logger.GetDefaultLogger(), types.Chains{
+	generator := NewReportGenerator(stateManager, logger.GetNopLogger(), types.Chains{
 		&types.Chain{Name: "chain"},
 	})
 
@@ -59,7 +44,7 @@ func TestReportGeneratorWithProposalError(t *testing.T) {
 func TestReportGeneratorWithVoteError(t *testing.T) {
 	t.Parallel()
 
-	stateManager := state.NewStateManager("./state.json", &TestFS{}, logger.GetDefaultLogger())
+	stateManager := state.NewStateManager("./state.json", &fs.TestFS{}, logger.GetNopLogger())
 
 	oldState := state.NewState()
 	newState := state.State{
@@ -83,7 +68,7 @@ func TestReportGeneratorWithVoteError(t *testing.T) {
 		},
 	}
 
-	generator := NewReportGenerator(stateManager, logger.GetDefaultLogger(), types.Chains{
+	generator := NewReportGenerator(stateManager, logger.GetNopLogger(), types.Chains{
 		&types.Chain{Name: "chain"},
 	})
 
@@ -98,7 +83,7 @@ func TestReportGeneratorWithVoteError(t *testing.T) {
 func TestReportGeneratorWithNotVoted(t *testing.T) {
 	t.Parallel()
 
-	stateManager := state.NewStateManager("./state.json", &TestFS{}, logger.GetDefaultLogger())
+	stateManager := state.NewStateManager("./state.json", &fs.TestFS{}, logger.GetNopLogger())
 
 	oldState := state.NewState()
 	newState := state.State{
@@ -118,7 +103,7 @@ func TestReportGeneratorWithNotVoted(t *testing.T) {
 		},
 	}
 
-	generator := NewReportGenerator(stateManager, logger.GetDefaultLogger(), types.Chains{
+	generator := NewReportGenerator(stateManager, logger.GetNopLogger(), types.Chains{
 		&types.Chain{Name: "chain"},
 	})
 
@@ -133,7 +118,7 @@ func TestReportGeneratorWithNotVoted(t *testing.T) {
 func TestReportGeneratorWithVoted(t *testing.T) {
 	t.Parallel()
 
-	stateManager := state.NewStateManager("./state.json", &TestFS{}, logger.GetDefaultLogger())
+	stateManager := state.NewStateManager("./state.json", &fs.TestFS{}, logger.GetNopLogger())
 
 	oldState := state.State{
 		ChainInfos: map[string]*state.ChainInfo{
@@ -172,7 +157,7 @@ func TestReportGeneratorWithVoted(t *testing.T) {
 		},
 	}
 
-	generator := NewReportGenerator(stateManager, logger.GetDefaultLogger(), types.Chains{
+	generator := NewReportGenerator(stateManager, logger.GetNopLogger(), types.Chains{
 		&types.Chain{Name: "chain"},
 	})
 
@@ -187,7 +172,7 @@ func TestReportGeneratorWithVoted(t *testing.T) {
 func TestReportGeneratorWithRevoted(t *testing.T) {
 	t.Parallel()
 
-	stateManager := state.NewStateManager("./state.json", &TestFS{}, logger.GetDefaultLogger())
+	stateManager := state.NewStateManager("./state.json", &fs.TestFS{}, logger.GetNopLogger())
 
 	oldState := state.State{
 		ChainInfos: map[string]*state.ChainInfo{
@@ -230,7 +215,7 @@ func TestReportGeneratorWithRevoted(t *testing.T) {
 		},
 	}
 
-	generator := NewReportGenerator(stateManager, logger.GetDefaultLogger(), types.Chains{
+	generator := NewReportGenerator(stateManager, logger.GetNopLogger(), types.Chains{
 		&types.Chain{Name: "chain"},
 	})
 

pkg/state/generator.go

@@ -1,23 +1,31 @@
 package state
 
 import (
-	"main/pkg/fetchers"
+	fetchersPkg "main/pkg/fetchers"
 	"main/pkg/types"
 	"sync"
 
 	"github.com/rs/zerolog"
 )
 
 type Generator struct {
-	Logger zerolog.Logger
-	Chains types.Chains
-	Mutex  sync.Mutex
+	Logger   zerolog.Logger
+	Chains   types.Chains
+	Fetchers map[string]fetchersPkg.Fetcher
+	Mutex    sync.Mutex
 }
 
 func NewStateGenerator(logger *zerolog.Logger, chains types.Chains) *Generator {
+	fetchers := make(map[string]fetchersPkg.Fetcher, len(chains))
+
+	for _, chain := range chains {
+		fetchers[chain.Name] = fetchersPkg.GetFetcher(chain, *logger)
+	}
+
 	return &Generator{
-		Logger: logger.With().Str("component", "state_generator").Logger(),
-		Chains: chains,
+		Logger:   logger.With().Str("component", "state_generator").Logger(),
+		Chains:   chains,
+		Fetchers: fetchers,
 	}
 }
 
@@ -29,8 +37,11 @@ func (g *Generator) GetState(oldState State) State {
 
 	for _, chain := range g.Chains {
 		g.Logger.Info().Str("name", chain.Name).Msg("Processing a chain")
+
+		fetcher := g.Fetchers[chain.Name]
+
 		go func(c *types.Chain) {
-			g.ProcessChain(c, state, oldState)
+			g.ProcessChain(c, state, oldState, fetcher)
 			wg.Done()
 		}(chain)
 	}
@@ -43,9 +54,8 @@ func (g *Generator) ProcessChain(
 	chain *types.Chain,
 	state State,
 	oldState State,
+	fetcher fetchersPkg.Fetcher,
 ) {
-	fetcher := fetchers.GetFetcher(chain, g.Logger)
-
 	prevHeight := oldState.GetLastProposalsHeight(chain)
 	proposals, proposalsHeight, err := fetcher.GetAllProposals(prevHeight)
 	if err != nil {
@@ -102,41 +112,45 @@ func (g *Generator) ProcessChain(
 func (g *Generator) ProcessProposalAndWallet(
 	chain *types.Chain,
 	proposal types.Proposal,
-	fetcher fetchers.Fetcher,
+	fetcher fetchersPkg.Fetcher,
 	wallet *types.Wallet,
 	state State,
 	oldState State,
 ) {
 	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
 	}