From 962c4902b6f0ef14886f725ec4352ec09b739460 Mon Sep 17 00:00:00 2001
From: Alex Gartner <git@agartner.com>
Date: Thu, 5 Dec 2024 13:41:33 -0800
Subject: [PATCH] add SubscribeNewBlocks test coverage

---
 zetaclient/testutils/mocks/cometbft_client.go | 32 ++++++++++++++++---
 zetaclient/zetacore/client_subscriptions.go   |  6 +---
 zetaclient/zetacore/client_test.go            | 30 +++++++++++++++++
 3 files changed, 59 insertions(+), 9 deletions(-)

diff --git a/zetaclient/testutils/mocks/cometbft_client.go b/zetaclient/testutils/mocks/cometbft_client.go
index dcb452621d..daa37a9467 100644
--- a/zetaclient/testutils/mocks/cometbft_client.go
+++ b/zetaclient/testutils/mocks/cometbft_client.go
@@ -3,6 +3,7 @@ package mocks
 import (
 	"context"
 	"encoding/hex"
+	"sync"
 	"testing"
 
 	abci "github.com/cometbft/cometbft/abci/types"
@@ -20,6 +21,9 @@ type CometBFTClient struct {
 	err    error
 	code   uint32
 	txHash bytes.HexBytes
+
+	subscribers     map[string]chan<- coretypes.ResultEvent
+	subscribersLock sync.Mutex
 }
 
 func (c *CometBFTClient) BroadcastTxCommit(
@@ -82,11 +86,31 @@ func (c *CometBFTClient) SetError(err error) *CometBFTClient {
 	return c
 }
 
+// PublishToSubscribers will publish an event to all subscribers (mock only)
+func (c *CometBFTClient) PublishToSubscribers(event coretypes.ResultEvent) {
+	c.subscribersLock.Lock()
+	defer c.subscribersLock.Unlock()
+	for _, ch := range c.subscribers {
+		ch <- event
+	}
+}
+
+func (c *CometBFTClient) Subscribe(
+	_ context.Context,
+	subscriber, _ string,
+	_ ...int,
+) (out <-chan coretypes.ResultEvent, err error) {
+	outChan := make(chan coretypes.ResultEvent)
+	c.subscribers[subscriber] = outChan
+	return outChan, nil
+}
+
 func NewSDKClientWithErr(t *testing.T, err error, code uint32) *CometBFTClient {
 	return &CometBFTClient{
-		t:      t,
-		Client: mock.Client{},
-		err:    err,
-		code:   code,
+		t:           t,
+		Client:      mock.Client{},
+		err:         err,
+		code:        code,
+		subscribers: make(map[string]chan<- coretypes.ResultEvent),
 	}
 }
diff --git a/zetaclient/zetacore/client_subscriptions.go b/zetaclient/zetacore/client_subscriptions.go
index 36f2aee2fa..364821efab 100644
--- a/zetaclient/zetacore/client_subscriptions.go
+++ b/zetaclient/zetacore/client_subscriptions.go
@@ -6,13 +6,9 @@ import (
 	cometbft_types "github.com/cometbft/cometbft/types"
 )
 
-const (
-	newBlockSubscriptionFilter = "tm.event='NewBlock'"
-)
-
 // NewBlockSubscriber subscribes to cometbft new block events
 func (c *Client) NewBlockSubscriber(ctx context.Context) (chan cometbft_types.EventDataNewBlock, error) {
-	rawBlockEventChan, err := c.cometBFTClient.Subscribe(ctx, "", newBlockSubscriptionFilter)
+	rawBlockEventChan, err := c.cometBFTClient.Subscribe(ctx, "", cometbft_types.EventQueryNewBlock.String())
 	if err != nil {
 		return nil, err
 	}
diff --git a/zetaclient/zetacore/client_test.go b/zetaclient/zetacore/client_test.go
index 940df90c0b..a445dcaf3c 100644
--- a/zetaclient/zetacore/client_test.go
+++ b/zetaclient/zetacore/client_test.go
@@ -22,6 +22,8 @@ import (
 	"go.nhat.io/grpcmock/planner"
 
 	cometbft_rpc_client "github.com/cometbft/cometbft/rpc/client"
+	coretypes "github.com/cometbft/cometbft/rpc/core/types"
+	cometbft_types "github.com/cometbft/cometbft/types"
 	"github.com/zeta-chain/node/cmd/zetacored/config"
 	crosschaintypes "github.com/zeta-chain/node/x/crosschain/types"
 	"github.com/zeta-chain/node/zetaclient/keys"
@@ -237,3 +239,31 @@ func TestZetacore_GetAllOutboundTrackerByChain(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, expectedOutput.OutboundTracker, resp)
 }
+
+func TestZetacore_SubscribeNewBlocks(t *testing.T) {
+	ctx := context.Background()
+	cometBFTClient := mocks.NewSDKClientWithErr(t, nil, 0)
+	client := setupZetacoreClient(
+		t,
+		withDefaultObserverKeys(),
+		withTendermint(cometBFTClient),
+	)
+
+	newBlockChan, err := client.NewBlockSubscriber(ctx)
+	require.NoError(t, err)
+
+	height := int64(10)
+
+	cometBFTClient.PublishToSubscribers(coretypes.ResultEvent{
+		Data: cometbft_types.EventDataNewBlock{
+			Block: &cometbft_types.Block{
+				Header: cometbft_types.Header{
+					Height: height,
+				},
+			},
+		},
+	})
+
+	newBlockEvent := <-newBlockChan
+	require.Equal(t, height, newBlockEvent.Block.Header.Height)
+}