fix flakey test

node/full_test.go

@@ -102,6 +102,9 @@ func TestFull_WithRemoteCore(t *testing.T) {
 
 	repo := MockRepository(t, DefaultConfig(Full))
 	remoteCore, protocol, ip := core.StartRemoteCore()
+	t.Cleanup(func() {
+		remoteCore.Stop() // nolint:errcheck
+	})
 	require.NotNil(t, remoteCore)
 	assert.True(t, remoteCore.IsRunning())
 

service/block/store_test.go

@@ -38,6 +38,9 @@ func TestService_BlockStore(t *testing.T) {
 
 func generateRawAndExtendedBlock(t *testing.T, store format.DAGService) *Block {
 	mockCore := core.MockEmbeddedClient()
+	t.Cleanup(func() {
+		mockCore.Stop() // nolint:errcheck
+	})
 	fetcher := core.NewBlockFetcher(mockCore)
 	sub, err := fetcher.SubscribeNewBlockEvent(context.Background())
 	require.NoError(t, err)

service/header/core_exchange_test.go

@@ -12,7 +12,11 @@ import (
 )
 
 func TestCoreExchange_RequestHeaders(t *testing.T) {
-	fetcher := createCoreFetcher()
+	mockClient, fetcher := createCoreFetcher()
+	t.Cleanup(func() {
+		mockClient.Stop() // nolint:errcheck
+	})
+
 	store := mdutils.Mock()
 
 	// generate 10 blocks
@@ -32,9 +36,9 @@ func Test_hashMatch(t *testing.T) {
 	assert.False(t, hashMatch(expected, mismatch))
 }
 
-func createCoreFetcher() *core.BlockFetcher {
+func createCoreFetcher() (core.Client, *core.BlockFetcher) {
 	mock := core.MockEmbeddedClient()
-	return core.NewBlockFetcher(mock)
+	return mock, core.NewBlockFetcher(mock)
 }
 
 func generateBlocks(t *testing.T, fetcher *core.BlockFetcher) {

service/header/core_listener_test.go

@@ -35,7 +35,10 @@ func TestCoreListener(t *testing.T) {
 	require.NoError(t, err)
 
 	md := mdutils.Mock()
-	fetcher := createCoreFetcher()
+	mockClient, fetcher := createCoreFetcher()
+	t.Cleanup(func() {
+		mockClient.Stop() // nolint:errcheck
+	})
 	ex := NewCoreExchange(fetcher, md)
 
 	cl, err := NewCoreListener(ex, topic1)
@@ -67,6 +70,10 @@ func TestCoreHeaderService(t *testing.T) {
 	// create a mockCore client that we can use later to generate fake blocks
 	coreNode := core.StartMockNode()
 	coreClient := core.NewEmbeddedFromNode(coreNode)
+	t.Cleanup(func() {
+		coreNode.Stop()   // nolint:errcheck
+		coreClient.Stop() // nolint:errcheck
+	})
 
 	fetcher := core.NewBlockFetcher(coreClient)
 	ex := NewCoreExchange(fetcher, md)
@@ -112,25 +119,32 @@ func TestCoreHeaderService(t *testing.T) {
 	// wait til syncing is finished to generate new blocks for the
 	// listener to catch and broadcast as ExtendedHeaders
 	<-syncer.done
+
 	// create a new subscription to generate another new block
 	blockGen := core.NewEmbeddedFromNode(coreNode)
 	blockSub, err := blockGen.Subscribe(context.Background(), "TEST",
 		types.QueryForEvent(types.EventNewBlock).String())
 	require.NoError(t, err)
-
-	<-blockSub
+	// get expected block details
+	rawExpected := <-blockSub
+	expected, ok := rawExpected.Data.(types.EventDataNewBlock)
+	if !ok {
+		t.Fatalf("unexpected: %v", rawExpected)
+	}
 	err = blockGen.Unsubscribe(context.Background(), "TEST",
 		types.QueryForEvent(types.EventNewBlock).String())
 	require.NoError(t, err)
 
-	// ensure block number 2 is broadcasted to gossipsub network
+	// ensure expected block is broadcasted to gossipsub network
 	msg, err := peerSub.Next(context.Background())
 	require.NoError(t, err)
 
+	// unmarshal gossipsub message into ExtendedHeader
 	var nextHeader ExtendedHeader
 	err = nextHeader.UnmarshalBinary(msg.Data)
 	require.NoError(t, err)
-	require.Equal(t, int64(2), nextHeader.Height)
+	require.Equal(t, expected.Block.Height, nextHeader.Height)
 
+	// stop header service
 	require.NoError(t, headerServ.Stop(context.Background()))
 }