add tests for aggregateMessages

Signed-off-by: Jun Kimura <[email protected]>
datachainlab · Jul 16, 2024 · 5c814d6 · 5c814d6
1 parent 512c9de
commit 5c814d6
Show file tree

Hide file tree

Showing 2 changed files with 111 additions and 6 deletions.
diff --git a/relay/prover.go b/relay/prover.go
@@ -193,7 +193,7 @@ func (pr *Prover) SetupHeadersForUpdate(dstChain core.FinalityAwareChain, latest
 	// NOTE: assume that the messages length and the signatures length are the same
 	if pr.config.MessageAggregation {
 		pr.getLogger().Info("aggregate messages", "num_messages", len(messages))
-		update, err := pr.aggregateMessages(messages, signatures, pr.activeEnclaveKey.EnclaveKeyAddress)
+		update, err := aggregateMessages(pr.getLogger(), pr.config.GetMessageAggregationBatchSize(), pr.lcpServiceClient.AggregateMessages, messages, signatures, pr.activeEnclaveKey.EnclaveKeyAddress)
 		if err != nil {
 			return nil, err
 		}
@@ -210,14 +210,23 @@ func (pr *Prover) SetupHeadersForUpdate(dstChain core.FinalityAwareChain, latest
 	return updates, nil
 }
 
-func (pr *Prover) aggregateMessages(messages [][]byte, signatures [][]byte, signer []byte) (*lcptypes.UpdateClientMessage, error) {
+type MessageAggregator func(ctx context.Context, in *elc.MsgAggregateMessages, opts ...grpc.CallOption) (*elc.MsgAggregateMessagesResponse, error)
+
+func aggregateMessages(
+	logger *log.RelayLogger,
+	batchSize uint64,
+	messageAggregator MessageAggregator,
+	messages [][]byte,
+	signatures [][]byte,
+	signer []byte,
+) (*lcptypes.UpdateClientMessage, error) {
 	if len(messages) == 0 {
 		return nil, fmt.Errorf("aggregateMessages: messages must not be empty")
 	} else if len(messages) != len(signatures) {
 		return nil, fmt.Errorf("aggregateMessages: messages and signatures must have the same length: messages=%v signatures=%v", len(messages), len(signatures))
 	}
 	for {
-		batches, err := splitIntoMultiBatch(messages, signatures, signer, pr.config.GetMessageAggregationBatchSize())
+		batches, err := splitIntoMultiBatch(messages, signatures, signer, batchSize)
 		if err != nil {
 			return nil, err
 		}
@@ -235,7 +244,7 @@ func (pr *Prover) aggregateMessages(messages [][]byte, signatures [][]byte, sign
 					Messages:   batches[0].Messages,
 					Signatures: batches[0].Signatures,
 				}
-				resp, err := pr.lcpServiceClient.AggregateMessages(context.TODO(), &m)
+				resp, err := messageAggregator(context.TODO(), &m)
 				if err != nil {
 					return nil, fmt.Errorf("failed to aggregate messages: msg=%v %w", m, err)
 				}
@@ -247,7 +256,7 @@ func (pr *Prover) aggregateMessages(messages [][]byte, signatures [][]byte, sign
 		} else if n == 0 {
 			return nil, fmt.Errorf("unexpected error: batches must not be empty")
 		} else {
-			pr.getLogger().Info("aggregateMessages", "num_batches", n)
+			logger.Info("aggregateMessages", "num_batches", n)
 		}
 		messages = nil
 		signatures = nil
@@ -257,7 +266,7 @@ func (pr *Prover) aggregateMessages(messages [][]byte, signatures [][]byte, sign
 				Messages:   b.Messages,
 				Signatures: b.Signatures,
 			}
-			resp, err := pr.lcpServiceClient.AggregateMessages(context.TODO(), &m)
+			resp, err := messageAggregator(context.TODO(), &m)
 			if err != nil {
 				return nil, fmt.Errorf("failed to aggregate messages: i=%v msg=%v %w", i, m, err)
 			}

diff --git a/relay/prover_test.go b/relay/prover_test.go
@@ -1,10 +1,14 @@
 package relay
 
 import (
+	"context"
 	"fmt"
 	"testing"
 
+	"github.com/datachainlab/lcp-go/relay/elc"
+	"github.com/hyperledger-labs/yui-relayer/log"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
 )
 
 func TestSplitIntoMultiBatch(t *testing.T) {
@@ -94,3 +98,95 @@ func TestSplitIntoMultiBatch(t *testing.T) {
 		})
 	}
 }
+
+func TestAggregateMessages(t *testing.T) {
+	var M = func(n uint64) []byte {
+		return []byte(fmt.Sprintf("message-%d", n))
+	}
+	var S = func(n uint64) []byte {
+		return []byte(fmt.Sprintf("signature-%d", n))
+	}
+	var Signer = func(n uint64) []byte {
+		return []byte(fmt.Sprintf("signer-%d", n))
+	}
+
+	err := log.InitLogger("DEBUG", "text", "stdout")
+	require.NoError(t, err)
+	logger := log.GetLogger()
+
+	var cases = []struct {
+		Messages   [][]byte
+		Signatures [][]byte
+		Signer     []byte
+		BatchSize  uint64
+		Error      bool
+	}{
+		// Messages.len = 0 is invalid
+		{
+			Messages:   [][]byte{},
+			Signatures: [][]byte{},
+			Signer:     Signer(0),
+			BatchSize:  2,
+			Error:      true,
+		},
+		{
+			Messages:   [][]byte{M(0)},
+			Signatures: [][]byte{S(0)},
+			Signer:     Signer(0),
+			BatchSize:  2,
+			Error:      false,
+		},
+		{
+			Messages:   [][]byte{M(0), M(1)},
+			Signatures: [][]byte{S(0), S(1)},
+			Signer:     Signer(0),
+			BatchSize:  2,
+			Error:      false,
+		},
+		// BatchSize = 1 is invalid
+		{
+			Messages:   [][]byte{M(0), M(1)},
+			Signatures: [][]byte{S(0), S(1)},
+			Signer:     Signer(0),
+			BatchSize:  1,
+			Error:      true,
+		},
+		{
+			Messages:   [][]byte{M(0), M(1), M(2)},
+			Signatures: [][]byte{S(0), S(1), S(2)},
+			Signer:     Signer(0),
+			BatchSize:  2,
+			Error:      false,
+		},
+	}
+
+	for i, c := range cases {
+		t.Run(fmt.Sprintf("case-%d", i), func(t *testing.T) {
+			require := require.New(t)
+			_, err := aggregateMessages(logger, c.BatchSize, mockMessageAggregator, c.Messages, c.Signatures, c.Signer)
+			if c.Error {
+				require.Error(err)
+				return
+			} else {
+				require.NoError(err)
+			}
+		})
+	}
+}
+
+func mockMessageAggregator(_ context.Context, in *elc.MsgAggregateMessages, _ ...grpc.CallOption) (*elc.MsgAggregateMessagesResponse, error) {
+	var res elc.MsgAggregateMessagesResponse
+	if len(in.Messages) != len(in.Signatures) {
+		return nil, fmt.Errorf("messages and signatures must have the same length")
+	}
+	if len(in.Messages) == 0 {
+		return nil, fmt.Errorf("messages.len = 0 is invalid")
+	} else if len(in.Messages) == 1 {
+		return nil, fmt.Errorf("messages.len = 1 is invalid")
+	}
+	for i := 0; i < len(in.Messages); i++ {
+		res.Message = append(res.Message, in.Messages[i]...)
+		res.Signature = append(res.Signature, in.Signatures[i]...)
+	}
+	return &res, nil
+}