From 07ab3d230025a7668f5d0d82fbdfacd0142e6375 Mon Sep 17 00:00:00 2001
From: Sam Liokumovich <65994425+samliok@users.noreply.github.com>
Date: Mon, 23 Dec 2024 15:54:35 -0500
Subject: [PATCH] cleanup epoch tests (#19)

* cleanup

* remove println

* add lock

* revert
---
 epoch.go                |  5 ++---
 epoch_multinode_test.go | 32 +++++++++++++++-----------------
 epoch_test.go           |  3 +++
 3 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/epoch.go b/epoch.go
index fbc1f2b..9d4473c 100644
--- a/epoch.go
+++ b/epoch.go
@@ -1039,8 +1039,7 @@ func (e *Epoch) storeProposal(block Block) bool {
 
 	// Have we already received a block from that node?
 	// If so, it cannot change its mind and send us a different block.
-	round, exists := e.rounds[md.Round]
-	if exists {
+	if _, exists := e.rounds[md.Round]; exists {
 		// We have already received a block for this round in the past, refuse receiving an alternative block.
 		// We do this because we may have already voted for a different block.
 		// Refuse processing the block to not be coerced into voting for a different block.
@@ -1048,7 +1047,7 @@ func (e *Epoch) storeProposal(block Block) bool {
 		return false
 	}
 
-	round = NewRound(block)
+	round := NewRound(block)
 	e.rounds[md.Round] = round
 	// We might have received votes and finalizations from future rounds before we received this block.
 	// So load the messages into our round data structure now that we have created it.
diff --git a/epoch_multinode_test.go b/epoch_multinode_test.go
index 91aa8c1..c93ed1d 100644
--- a/epoch_multinode_test.go
+++ b/epoch_multinode_test.go
@@ -6,10 +6,11 @@ package simplex_test
 import (
 	"bytes"
 	"context"
-	"github.com/stretchr/testify/require"
 	. "simplex"
 	"simplex/wal"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestSimplexMultiNodeSimple(t *testing.T) {
@@ -29,14 +30,19 @@ func TestSimplexMultiNodeSimple(t *testing.T) {
 		n.start()
 	}
 
-	for seq := 0; seq < 1; seq++ {
+	for seq := 0; seq < 10; seq++ {
 		for _, n := range instances {
 			n.ledger.waitForBlockCommit(uint64(seq))
-			bb.triggerNewBlock()
 		}
+		bb.triggerNewBlock()
 	}
 }
 
+func (t *testInstance) start() {
+	require.NoError(t.t, t.e.Start())
+	go t.handleMessages()
+}
+
 func newSimplexNode(t *testing.T, id uint8, net *inMemNetwork, bb BlockBuilder) *testInstance {
 	l := makeLogger(t, int(id))
 	storage := newInMemStorage()
@@ -84,20 +90,12 @@ type testInstance struct {
 	t *testing.T
 }
 
-func (t *testInstance) start() {
-	require.NoError(t.t, t.e.Start())
-	go t.run()
-}
-
-func (t *testInstance) run() {
-	for {
-		select {
-		case msg := <-t.ingress:
-			err := t.e.HandleMessage(msg.msg, msg.from)
-			require.NoError(t.t, err)
-			if err != nil {
-				return
-			}
+func (t *testInstance) handleMessages() {
+	for msg := range t.ingress {
+		err := t.e.HandleMessage(msg.msg, msg.from)
+		require.NoError(t.t, err)
+		if err != nil {
+			return
 		}
 	}
 }
diff --git a/epoch_test.go b/epoch_test.go
index 8f2dd99..23e4113 100644
--- a/epoch_test.go
+++ b/epoch_test.go
@@ -278,6 +278,9 @@ func (mem *InMemStorage) Retrieve(seq uint64) (Block, FinalizationCertificate, b
 }
 
 func (mem *InMemStorage) Index(seq uint64, block Block, certificate FinalizationCertificate) {
+	mem.lock.Lock()
+	defer mem.lock.Unlock()
+	
 	_, ok := mem.data[seq]
 	if ok {
 		panic(fmt.Sprintf("block with seq %d already indexed!", seq))