Add unit test for restore execution

backup/backup.go

@@ -1,6 +1,7 @@
 package backup
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -12,6 +13,7 @@ import (
 
 // ExecuteBackup executes the node backup process
 func ExecuteBackup(
+	ctx context.Context,
 	client client.Client,
 	writer io.Writer,
 	logger log.Logger,
@@ -30,26 +32,33 @@ func ExecuteBackup(
 
 	// Gather the chain data from the node
 	for block := cfg.FromBlock; block <= toBlock; block++ {
-		txs, txErr := client.GetBlockTransactions(block)
-		if txErr != nil {
-			return fmt.Errorf("unable to fetch block transactions, %w", txErr)
-		}
-
-		// Save the block transaction data, if any
-		for _, tx := range txs {
-			data := &types.TxData{
-				Tx:       tx,
-				BlockNum: block,
+		select {
+		case <-ctx.Done():
+			logger.Info("export procedure stopped")
+
+			return nil
+		default:
+			txs, txErr := client.GetBlockTransactions(block)
+			if txErr != nil {
+				return fmt.Errorf("unable to fetch block transactions, %w", txErr)
 			}
 
-			// Write the tx data to the file
-			if writeErr := writeTxData(writer, data); writeErr != nil {
-				return fmt.Errorf("unable to write tx data, %w", writeErr)
+			// Save the block transaction data, if any
+			for _, tx := range txs {
+				data := &types.TxData{
+					Tx:       tx,
+					BlockNum: block,
+				}
+
+				// Write the tx data to the file
+				if writeErr := writeTxData(writer, data); writeErr != nil {
+					return fmt.Errorf("unable to write tx data, %w", writeErr)
+				}
 			}
-		}
 
-		// Log the progress
-		logProgress(logger, cfg.FromBlock, toBlock, block)
+			// Log the progress
+			logProgress(logger, cfg.FromBlock, toBlock, block)
+		}
 	}
 
 	return nil

backup/backup_test.go

@@ -2,6 +2,7 @@ package backup
 
 import (
 	"bufio"
+	"context"
 	"encoding/json"
 	"errors"
 	"os"
@@ -119,7 +120,16 @@ func TestBackup_ExecuteBackup(t *testing.T) {
 	cfg.ToBlock = &toBlock
 
 	// Run the backup procedure
-	require.NoError(t, ExecuteBackup(mockClient, tempFile, noop.New(), cfg))
+	require.NoError(
+		t,
+		ExecuteBackup(
+			context.Background(),
+			mockClient,
+			tempFile,
+			noop.New(),
+			cfg,
+		),
+	)
 
 	// Read the output file
 	fileRaw, err := os.Open(tempFile.Name())

restore/mock_test.go

@@ -0,0 +1,49 @@
+package restore
+
+import (
+	"context"
+
+	"github.com/gnolang/gno/tm2/pkg/std"
+)
+
+type (
+	sendTransactionDelegate func(*std.Tx) error
+)
+
+type mockClient struct {
+	sendTransactionFn sendTransactionDelegate
+}
+
+func (m *mockClient) SendTransaction(tx *std.Tx) error {
+	if m.sendTransactionFn != nil {
+		return m.sendTransactionFn(tx)
+	}
+
+	return nil
+}
+
+type (
+	nextDelegate  func(context.Context) (*std.Tx, error)
+	closeDelegate func() error
+)
+
+type mockSource struct {
+	nextFn  nextDelegate
+	closeFn closeDelegate
+}
+
+func (m *mockSource) Next(ctx context.Context) (*std.Tx, error) {
+	if m.nextFn != nil {
+		return m.nextFn(ctx)
+	}
+
+	return nil, nil
+}
+
+func (m *mockSource) Close() error {
+	if m.closeFn != nil {
+		return m.closeFn()
+	}
+
+	return nil
+}

restore/restore.go

@@ -14,25 +14,23 @@ import (
 
 // ExecuteRestore executes the node restore process
 func ExecuteRestore(
+	ctx context.Context,
 	client client.Client,
 	source source.Source,
 	logger log.Logger,
-	cfg Config,
 ) error {
-	// Verify the config
-	if cfgErr := ValidateConfig(cfg); cfgErr != nil {
-		return fmt.Errorf("invalid config, %w", cfgErr)
-	}
-
-	defer func() {
+	// Set up the teardown
+	teardown := func() {
 		if closeErr := source.Close(); closeErr != nil {
 			logger.Error(
 				"unable to gracefully close source",
 				"err",
 				closeErr.Error(),
 			)
 		}
-	}()
+	}
+
+	defer teardown()
 
 	var (
 		tx      *std.Tx
@@ -42,8 +40,12 @@ func ExecuteRestore(
 	)
 
 	// Fetch next transactions
-	// TODO add ctx
-	for tx, nextErr = source.Next(context.Background()); nextErr == nil; {
+	for nextErr == nil {
+		tx, nextErr = source.Next(ctx)
+		if nextErr != nil {
+			break
+		}
+
 		// Send the transaction
 		if sendErr := client.SendTransaction(tx); sendErr != nil {
 			// Invalid transaction sends are only logged,
@@ -61,16 +63,16 @@ func ExecuteRestore(
 	}
 
 	// Check if this is the end of the road
-	if errors.Is(nextErr, io.EOF) {
-		// No more transactions to apply
-		logger.Info(
-			"restore process finished",
-			"total",
-			totalTxs,
-		)
-
-		return nil
+	if !errors.Is(nextErr, io.EOF) {
+		return fmt.Errorf("unable to get next transaction, %w", nextErr)
 	}
 
-	return fmt.Errorf("unable to get next transaction, %w", nextErr)
+	// No more transactions to apply
+	logger.Info(
+		"restore process finished",
+		"total",
+		totalTxs,
+	)
+
+	return nil
 }

restore/restore_test.go

@@ -0,0 +1,63 @@
+package restore
+
+import (
+	"context"
+	"io"
+	"testing"
+
+	"github.com/gnolang/gno/tm2/pkg/std"
+	"github.com/gnolang/tx-archive/log/noop"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRestore_ExecuteRestore(t *testing.T) {
+	t.Parallel()
+
+	var (
+		exampleTxCount = 10
+		exampleTxGiven = 0
+
+		exampleTx = &std.Tx{
+			Memo: "example tx",
+		}
+
+		sentTxs = make([]*std.Tx, 0)
+
+		mockClient = &mockClient{
+			sendTransactionFn: func(tx *std.Tx) error {
+				sentTxs = append(sentTxs, tx)
+
+				return nil
+			},
+		}
+		mockSource = &mockSource{
+			nextFn: func(ctx context.Context) (*std.Tx, error) {
+				if exampleTxGiven == exampleTxCount {
+					return nil, io.EOF
+				}
+
+				exampleTxGiven++
+
+				return exampleTx, nil
+			},
+		}
+	)
+
+	// Execute the restore
+	assert.NoError(
+		t,
+		ExecuteRestore(
+			context.Background(),
+			mockClient,
+			mockSource,
+			noop.New(),
+		),
+	)
+
+	// Verify the restore was correct
+	assert.Len(t, sentTxs, exampleTxCount)
+
+	for _, tx := range sentTxs {
+		assert.Equal(t, exampleTx, tx)
+	}
+}