From a665e679db2bbbebc121967c8221dc61ed6294eb Mon Sep 17 00:00:00 2001
From: Dmitry <dmitry@zetachain.com>
Date: Fri, 2 Aug 2024 16:31:36 +0200
Subject: [PATCH] Fix merge conflicts

---
 zetaclient/maintenance/tss_listener.go        |  4 +-
 zetaclient/testutils/mocks/zetacore_client.go | 58 +++++++++++++++++++
 2 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/zetaclient/maintenance/tss_listener.go b/zetaclient/maintenance/tss_listener.go
index bb7b1fb050..24c0f12b8b 100644
--- a/zetaclient/maintenance/tss_listener.go
+++ b/zetaclient/maintenance/tss_listener.go
@@ -129,7 +129,7 @@ func (tl *TSSListener) waitForNewKeyGeneration(ctx context.Context) error {
 func (tl *TSSListener) waitForNewKeygen(ctx context.Context) error {
 	// Initial Keygen retrieval
 	keygen, err := retry.DoTypedWithBackoffAndRetry(
-		func() (*observertypes.Keygen, error) { return tl.client.GetKeyGen(ctx) },
+		func() (observertypes.Keygen, error) { return tl.client.GetKeyGen(ctx) },
 		retry.DefaultConstantBackoff(),
 	)
 	if err != nil {
@@ -147,8 +147,6 @@ func (tl *TSSListener) waitForNewKeygen(ctx context.Context) error {
 			case err != nil:
 				tl.logger.Warn().Err(err).Msg("unable to get keygen")
 				continue
-			case keygenUpdated == nil:
-				continue
 			case keygenUpdated.Status == observertypes.KeygenStatus_PendingKeygen:
 				continue
 			case keygen.BlockNumber == keygenUpdated.BlockNumber:
diff --git a/zetaclient/testutils/mocks/zetacore_client.go b/zetaclient/testutils/mocks/zetacore_client.go
index 168b580ada..5efe76aeba 100644
--- a/zetaclient/testutils/mocks/zetacore_client.go
+++ b/zetaclient/testutils/mocks/zetacore_client.go
@@ -496,6 +496,64 @@ func (_m *ZetacoreClient) GetRateLimiterInput(ctx context.Context, window int64)
 	return r0, r1
 }
 
+// GetTSS provides a mock function with given fields: ctx
+func (_m *ZetacoreClient) GetTSS(ctx context.Context) (observertypes.TSS, error) {
+	ret := _m.Called(ctx)
+
+	if len(ret) == 0 {
+		panic("no return value specified for GetTSS")
+	}
+
+	var r0 observertypes.TSS
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context) (observertypes.TSS, error)); ok {
+		return rf(ctx)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context) observertypes.TSS); ok {
+		r0 = rf(ctx)
+	} else {
+		r0 = ret.Get(0).(observertypes.TSS)
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context) error); ok {
+		r1 = rf(ctx)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// GetTSSHistory provides a mock function with given fields: ctx
+func (_m *ZetacoreClient) GetTSSHistory(ctx context.Context) ([]observertypes.TSS, error) {
+	ret := _m.Called(ctx)
+
+	if len(ret) == 0 {
+		panic("no return value specified for GetTSSHistory")
+	}
+
+	var r0 []observertypes.TSS
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context) ([]observertypes.TSS, error)); ok {
+		return rf(ctx)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context) []observertypes.TSS); ok {
+		r0 = rf(ctx)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]observertypes.TSS)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context) error); ok {
+		r1 = rf(ctx)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
 // GetZetaHotKeyBalance provides a mock function with given fields: ctx
 func (_m *ZetacoreClient) GetZetaHotKeyBalance(ctx context.Context) (math.Int, error) {
 	ret := _m.Called(ctx)