From 25afe20fea420b930a8a221cc4c78620b0a7b510 Mon Sep 17 00:00:00 2001
From: Michael Tsitrin <114929630+mtsitrin@users.noreply.github.com>
Date: Wed, 17 Apr 2024 11:11:01 +0300
Subject: [PATCH] feat(perf): removed unneeded state-index query (#650)

---
 settlement/dymension/dymension.go      | 16 +++++-----------
 settlement/dymension/dymension_test.go |  5 +----
 2 files changed, 6 insertions(+), 15 deletions(-)

diff --git a/settlement/dymension/dymension.go b/settlement/dymension/dymension.go
index 32f7aacf0..7b7148905 100644
--- a/settlement/dymension/dymension.go
+++ b/settlement/dymension/dymension.go
@@ -286,12 +286,11 @@ func (d *HubClient) PostBatch(batch *types.Batch, daClient da.Client, daResult *
 
 // GetLatestBatch returns the latest batch from the Dymension Hub.
 func (d *HubClient) GetLatestBatch(rollappID string) (*settlement.ResultRetrieveBatch, error) {
-	var latestStateInfoIndexResp *rollapptypes.QueryGetLatestStateIndexResponse
-
+	var stateInfoResp *rollapptypes.QueryGetStateInfoResponse
 	err := d.RunWithRetry(func() error {
 		var err error
-		latestStateInfoIndexResp, err = d.rollappQueryClient.LatestStateIndex(d.ctx,
-			&rollapptypes.QueryGetLatestStateIndexRequest{RollappId: d.config.RollappID})
+		stateInfoResp, err = d.rollappQueryClient.StateInfo(d.ctx,
+			&rollapptypes.QueryGetStateInfoRequest{RollappId: d.config.RollappID})
 
 		if status.Code(err) == codes.NotFound {
 			return retry.Unrecoverable(settlement.ErrBatchNotFound)
@@ -302,17 +301,12 @@ func (d *HubClient) GetLatestBatch(rollappID string) (*settlement.ResultRetrieve
 	if err != nil {
 		return nil, err
 	}
-
 	// not supposed to happen, but just in case
-	if latestStateInfoIndexResp == nil {
+	if stateInfoResp == nil {
 		return nil, settlement.ErrEmptyResponse
 	}
 
-	latestBatch, err := d.GetBatchAtIndex(rollappID, latestStateInfoIndexResp.StateIndex.Index)
-	if err != nil {
-		return nil, err
-	}
-	return latestBatch, nil
+	return d.convertStateInfoToResultRetrieveBatch(&stateInfoResp.StateInfo)
 }
 
 // GetBatchAtIndex returns the batch at the given index from the Dymension Hub.
diff --git a/settlement/dymension/dymension_test.go b/settlement/dymension/dymension_test.go
index 783a0ccb7..f619e77fe 100644
--- a/settlement/dymension/dymension_test.go
+++ b/settlement/dymension/dymension_test.go
@@ -164,7 +164,6 @@ func TestPostBatch(t *testing.T) {
 			wg.Add(eventsCount)
 			// Reset the mock functions
 			testutil.UnsetMockFn(cosmosClientMock.On("BroadcastTx"))
-			testutil.UnsetMockFn(rollappQueryClientMock.On("LatestStateIndex"))
 			testutil.UnsetMockFn(rollappQueryClientMock.On("StateInfo"))
 			// Set the mock logic based on the test case
 			if !c.isBatchSubmitSuccess {
@@ -174,8 +173,6 @@ func TestPostBatch(t *testing.T) {
 			}
 			if c.shouldMockBatchIncluded {
 				if c.isBatchIncludedSuccess {
-					rollappQueryClientMock.On("LatestStateIndex", mock.Anything, mock.Anything).Return(
-						&rollapptypes.QueryGetLatestStateIndexResponse{StateIndex: rollapptypes.StateInfoIndex{Index: 1}}, nil)
 					daMetaData := &da.DASubmitMetaData{
 						Height: 1,
 						Client: da.Mock,
@@ -186,7 +183,7 @@ func TestPostBatch(t *testing.T) {
 						}},
 						nil)
 				} else {
-					rollappQueryClientMock.On("LatestStateIndex", mock.Anything, mock.Anything).Return(nil, fmt.Errorf("error"))
+					rollappQueryClientMock.On("StateInfo", mock.Anything, mock.Anything).Return(nil, fmt.Errorf("error"))
 				}
 			}
 			hubClient, err := newDymensionHubClient(settlement.Config{}, pubsubServer, log.TestingLogger(), options...)