From eba17dfbca374640c8275c0873aab9dbdf980df1 Mon Sep 17 00:00:00 2001
From: Alex Gartner <git@agartner.com>
Date: Mon, 16 Dec 2024 13:38:50 -0800
Subject: [PATCH] zetaclient test coverage

---
 zetaclient/maintenance/shutdown_listener.go   |   2 +-
 .../maintenance/shutdown_listener_test.go     | 105 ++++++++++++++++++
 2 files changed, 106 insertions(+), 1 deletion(-)
 create mode 100644 zetaclient/maintenance/shutdown_listener_test.go

diff --git a/zetaclient/maintenance/shutdown_listener.go b/zetaclient/maintenance/shutdown_listener.go
index 234d11cbaa..0ac3b18d75 100644
--- a/zetaclient/maintenance/shutdown_listener.go
+++ b/zetaclient/maintenance/shutdown_listener.go
@@ -86,7 +86,7 @@ func (o *ShutdownListener) handleNewFlags(ctx context.Context, f observertypes.O
 		return false
 	}
 
-	if f.RestartHeight > currentHeight {
+	if f.RestartHeight < currentHeight {
 		// only log restart height misseed once
 		if o.lastRestartHeightMissed != f.RestartHeight {
 			o.logger.Error().
diff --git a/zetaclient/maintenance/shutdown_listener_test.go b/zetaclient/maintenance/shutdown_listener_test.go
new file mode 100644
index 0000000000..47ee8c7b65
--- /dev/null
+++ b/zetaclient/maintenance/shutdown_listener_test.go
@@ -0,0 +1,105 @@
+package maintenance
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	cometbfttypes "github.com/cometbft/cometbft/types"
+	"github.com/rs/zerolog"
+	"github.com/stretchr/testify/require"
+	observertypes "github.com/zeta-chain/node/x/observer/types"
+	"github.com/zeta-chain/node/zetaclient/testutils/mocks"
+)
+
+func assertChannelNotClosed[T any](t *testing.T, ch <-chan T) {
+	select {
+	case <-ch:
+		t.FailNow()
+	default:
+	}
+}
+
+func newBlockEventHeightOnly(height int64) cometbfttypes.EventDataNewBlock {
+	return cometbfttypes.EventDataNewBlock{
+		Block: &cometbfttypes.Block{
+			Header: cometbfttypes.Header{
+				Height: height,
+			},
+		},
+	}
+}
+
+func TestShutdownListener(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	logger := zerolog.New(zerolog.NewTestWriter(t))
+
+	t.Run("scheduled shutdown", func(t *testing.T) {
+		client := mocks.NewZetacoreClient(t)
+
+		listener := NewShutdownListener(client, logger)
+
+		client.Mock.On("GetOperationalFlags", ctx).Return(observertypes.OperationalFlags{
+			RestartHeight: 10,
+		}, nil)
+		client.Mock.On("GetBlockHeight", ctx).Return(int64(8), nil)
+		blockChan := make(chan cometbfttypes.EventDataNewBlock)
+		client.Mock.On("NewBlockSubscriber", ctx).Return(blockChan, nil)
+
+		complete := make(chan interface{})
+		listener.Listen(ctx, func() {
+			close(complete)
+		})
+
+		assertChannelNotClosed(t, complete)
+
+		blockChan <- newBlockEventHeightOnly(9)
+		assertChannelNotClosed(t, complete)
+
+		blockChan <- newBlockEventHeightOnly(10)
+		<-complete
+	})
+
+	t.Run("no shutdown", func(t *testing.T) {
+		client := mocks.NewZetacoreClient(t)
+
+		listener := NewShutdownListener(client, logger)
+
+		client.Mock.On("GetOperationalFlags", ctx).Return(observertypes.OperationalFlags{}, nil)
+		// GetBlockHeight is not mocked because we want the test to panic if it's called
+		// NewBlockSubscriber is not mocked because we want the test to panic if it's called
+		complete := make(chan interface{})
+		listener.Listen(ctx, func() {
+			close(complete)
+		})
+
+		require.Eventually(t, func() bool {
+			return len(client.Calls) == 1
+		}, time.Second, time.Millisecond)
+
+		assertChannelNotClosed(t, complete)
+	})
+
+	t.Run("shutdown height missed", func(t *testing.T) {
+		client := mocks.NewZetacoreClient(t)
+
+		listener := NewShutdownListener(client, logger)
+
+		client.Mock.On("GetOperationalFlags", ctx).Return(observertypes.OperationalFlags{
+			RestartHeight: 10,
+		}, nil)
+		client.Mock.On("GetBlockHeight", ctx).Return(int64(11), nil)
+
+		complete := make(chan interface{})
+		listener.Listen(ctx, func() {
+			close(complete)
+		})
+
+		require.Eventually(t, func() bool {
+			return len(client.Calls) == 2
+		}, time.Second, time.Millisecond)
+		assertChannelNotClosed(t, complete)
+
+	})
+}