From b1c090d12ca304e12436364af74b6791e6aa2786 Mon Sep 17 00:00:00 2001
From: Andrii <andriy.dyachuk95@gmail.com>
Date: Wed, 18 Sep 2024 12:44:28 +0300
Subject: [PATCH] Added new endpoints for getting state snapshot

---
 access/grpc/client.go    |  8 ++++++
 access/grpc/grpc.go      | 26 ++++++++++++++++++
 access/grpc/grpc_test.go | 58 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 92 insertions(+)

diff --git a/access/grpc/client.go b/access/grpc/client.go
index 5ca1729aa..b99ad64e0 100644
--- a/access/grpc/client.go
+++ b/access/grpc/client.go
@@ -236,6 +236,14 @@ func (c *Client) GetLatestProtocolStateSnapshot(ctx context.Context) ([]byte, er
 	return c.grpc.GetLatestProtocolStateSnapshot(ctx)
 }
 
+func (c *Client) GetProtocolStateSnapshotByBlockID(ctx context.Context, blockID flow.Identifier) ([]byte, error) {
+	return c.grpc.GetProtocolStateSnapshotByBlockID(ctx, blockID)
+}
+
+func (c *Client) GetProtocolStateSnapshotByHeight(ctx context.Context, blockHeight uint64) ([]byte, error) {
+	return c.grpc.GetProtocolStateSnapshotByHeight(ctx, blockHeight)
+}
+
 func (c *Client) GetExecutionResultForBlockID(ctx context.Context, blockID flow.Identifier) (*flow.ExecutionResult, error) {
 	return c.grpc.GetExecutionResultForBlockID(ctx, blockID)
 }
diff --git a/access/grpc/grpc.go b/access/grpc/grpc.go
index b70a55bed..11a408164 100644
--- a/access/grpc/grpc.go
+++ b/access/grpc/grpc.go
@@ -729,6 +729,32 @@ func (c *BaseClient) GetLatestProtocolStateSnapshot(ctx context.Context, opts ..
 	return res.GetSerializedSnapshot(), nil
 }
 
+func (c *BaseClient) GetProtocolStateSnapshotByBlockID(ctx context.Context, blockID flow.Identifier, opts ...grpc.CallOption) ([]byte, error) {
+	req := &access.GetProtocolStateSnapshotByBlockIDRequest{
+		BlockId: blockID.Bytes(),
+	}
+
+	res, err := c.rpcClient.GetProtocolStateSnapshotByBlockID(ctx, req, opts...)
+	if err != nil {
+		return nil, newRPCError(err)
+	}
+
+	return res.GetSerializedSnapshot(), nil
+}
+
+func (c *BaseClient) GetProtocolStateSnapshotByHeight(ctx context.Context, blockHeight uint64, opts ...grpc.CallOption) ([]byte, error) {
+	req := &access.GetProtocolStateSnapshotByHeightRequest{
+		BlockHeight: blockHeight,
+	}
+
+	res, err := c.rpcClient.GetProtocolStateSnapshotByHeight(ctx, req, opts...)
+	if err != nil {
+		return nil, newRPCError(err)
+	}
+
+	return res.GetSerializedSnapshot(), nil
+}
+
 func (c *BaseClient) GetExecutionResultForBlockID(ctx context.Context, blockID flow.Identifier, opts ...grpc.CallOption) (*flow.ExecutionResult, error) {
 	er, err := c.rpcClient.GetExecutionResultForBlockID(ctx, &access.GetExecutionResultForBlockIDRequest{
 		BlockId: convert.IdentifierToMessage(blockID),
diff --git a/access/grpc/grpc_test.go b/access/grpc/grpc_test.go
index 58a7939f6..655bdb3e4 100644
--- a/access/grpc/grpc_test.go
+++ b/access/grpc/grpc_test.go
@@ -1228,6 +1228,64 @@ func TestClient_GetLatestProtocolStateSnapshot(t *testing.T) {
 	}))
 }
 
+func TestClient_GetProtocolStateSnapshotByBlockID(t *testing.T) {
+	ids := test.IdentifierGenerator()
+
+	t.Run("Success", clientTest(func(t *testing.T, ctx context.Context, rpc *mocks.MockRPCClient, c *BaseClient) {
+		blockID := ids.New()
+
+		expected := &access.ProtocolStateSnapshotResponse{
+			SerializedSnapshot: make([]byte, 128),
+		}
+		_, err := rand.Read(expected.SerializedSnapshot)
+		assert.NoError(t, err)
+
+		rpc.On("GetProtocolStateSnapshotByBlockID", ctx, mock.Anything).Return(expected, nil)
+
+		res, err := c.GetProtocolStateSnapshotByBlockID(ctx, blockID)
+		assert.NoError(t, err)
+		assert.Equal(t, expected.SerializedSnapshot, res)
+	}))
+
+	t.Run("Internal error", clientTest(func(t *testing.T, ctx context.Context, rpc *mocks.MockRPCClient, c *BaseClient) {
+		blockID := ids.New()
+
+		rpc.On("GetProtocolStateSnapshotByBlockID", ctx, mock.Anything).
+			Return(nil, errInternal)
+
+		_, err := c.GetProtocolStateSnapshotByBlockID(ctx, blockID)
+		assert.Error(t, err)
+		assert.Equal(t, codes.Internal, status.Code(err))
+	}))
+}
+
+func TestClient_GetProtocolStateSnapshotByHeight(t *testing.T) {
+	blockHeight := uint64(42)
+
+	t.Run("Success", clientTest(func(t *testing.T, ctx context.Context, rpc *mocks.MockRPCClient, c *BaseClient) {
+		expected := &access.ProtocolStateSnapshotResponse{
+			SerializedSnapshot: make([]byte, 128),
+		}
+		_, err := rand.Read(expected.SerializedSnapshot)
+		assert.NoError(t, err)
+
+		rpc.On("GetProtocolStateSnapshotByHeight", ctx, mock.Anything).Return(expected, nil)
+
+		res, err := c.GetProtocolStateSnapshotByHeight(ctx, blockHeight)
+		assert.NoError(t, err)
+		assert.Equal(t, expected.SerializedSnapshot, res)
+	}))
+
+	t.Run("Internal error", clientTest(func(t *testing.T, ctx context.Context, rpc *mocks.MockRPCClient, c *BaseClient) {
+		rpc.On("GetProtocolStateSnapshotByHeight", ctx, mock.Anything).
+			Return(nil, errInternal)
+
+		_, err := c.GetProtocolStateSnapshotByHeight(ctx, blockHeight)
+		assert.Error(t, err)
+		assert.Equal(t, codes.Internal, status.Code(err))
+	}))
+}
+
 func TestClient_GetExecutionResultForBlockID(t *testing.T) {
 	ids := test.IdentifierGenerator()
 	t.Run("Success", clientTest(func(t *testing.T, ctx context.Context, rpc *mocks.MockRPCClient, c *BaseClient) {