github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/cmd/util/ledger/migrations/account_based_migration_test.go (about)

     1  package migrations
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  
     7  	"testing"
     8  
     9  	"github.com/onflow/cadence/runtime/common"
    10  	"github.com/rs/zerolog"
    11  	"github.com/stretchr/testify/require"
    12  
    13  	"github.com/onflow/flow-go/cmd/util/ledger/util/registers"
    14  	"github.com/onflow/flow-go/fvm/environment"
    15  	"github.com/onflow/flow-go/ledger"
    16  	"github.com/onflow/flow-go/ledger/common/convert"
    17  	"github.com/onflow/flow-go/model/flow"
    18  )
    19  
    20  func accountStatusPayload(address common.Address) *ledger.Payload {
    21  	accountStatus := environment.NewAccountStatus()
    22  
    23  	return ledger.NewPayload(
    24  		convert.RegisterIDToLedgerKey(
    25  			flow.AccountStatusRegisterID(flow.ConvertAddress(address)),
    26  		),
    27  		accountStatus.ToBytes(),
    28  	)
    29  }
    30  
    31  func TestErrorPropagation(t *testing.T) {
    32  	t.Parallel()
    33  
    34  	log := zerolog.New(zerolog.NewTestWriter(t))
    35  
    36  	address, err := common.HexToAddress("0x1")
    37  	require.NoError(t, err)
    38  
    39  	migrateWith := func(mig AccountBasedMigration) error {
    40  
    41  		// at least one payload otherwise the migration will not get called
    42  		payloads := []*ledger.Payload{
    43  			accountStatusPayload(address),
    44  		}
    45  
    46  		registersByAccount, err := registers.NewByAccountFromPayloads(payloads)
    47  		if err != nil {
    48  			return fmt.Errorf("could not create registers by account: %w", err)
    49  		}
    50  
    51  		return MigrateByAccount(
    52  			log,
    53  			10,
    54  			registersByAccount,
    55  			[]AccountBasedMigration{
    56  				mig,
    57  			},
    58  		)
    59  	}
    60  
    61  	t.Run("no err", func(t *testing.T) {
    62  		t.Parallel()
    63  
    64  		err := migrateWith(
    65  			testMigration{},
    66  		)
    67  		require.NoError(t, err)
    68  	})
    69  
    70  	t.Run("err on close", func(t *testing.T) {
    71  		t.Parallel()
    72  
    73  		desiredErr := fmt.Errorf("test close error")
    74  		err := migrateWith(
    75  			testMigration{
    76  				CloseFN: func() error {
    77  					return desiredErr
    78  				},
    79  			},
    80  		)
    81  		require.ErrorIs(t, err, desiredErr)
    82  	})
    83  
    84  	t.Run("err on init", func(t *testing.T) {
    85  		t.Parallel()
    86  
    87  		desiredErr := fmt.Errorf("test init error")
    88  		err := migrateWith(
    89  			testMigration{
    90  				InitMigrationFN: func(
    91  					log zerolog.Logger,
    92  					registersByAccount *registers.ByAccount,
    93  					nWorkers int,
    94  				) error {
    95  					return desiredErr
    96  				},
    97  			},
    98  		)
    99  		require.ErrorIs(t, err, desiredErr)
   100  	})
   101  
   102  	t.Run("err on migrate", func(t *testing.T) {
   103  		t.Parallel()
   104  
   105  		desiredErr := fmt.Errorf("test migrate error")
   106  		err := migrateWith(
   107  			testMigration{
   108  				MigrateAccountFN: func(
   109  					_ context.Context,
   110  					_ common.Address,
   111  					_ *registers.AccountRegisters,
   112  				) error {
   113  					return desiredErr
   114  				},
   115  			},
   116  		)
   117  		require.ErrorIs(t, err, desiredErr)
   118  	})
   119  }
   120  
   121  type testMigration struct {
   122  	InitMigrationFN func(
   123  		log zerolog.Logger,
   124  		registersByAccount *registers.ByAccount,
   125  		nWorkers int,
   126  	) error
   127  	MigrateAccountFN func(
   128  		ctx context.Context,
   129  		address common.Address,
   130  		accountRegisters *registers.AccountRegisters,
   131  	) error
   132  	CloseFN func() error
   133  }
   134  
   135  var _ AccountBasedMigration = &testMigration{}
   136  
   137  func (t testMigration) InitMigration(
   138  	log zerolog.Logger,
   139  	registersByAccount *registers.ByAccount,
   140  	nWorkers int,
   141  ) error {
   142  	if t.InitMigrationFN != nil {
   143  		return t.InitMigrationFN(log, registersByAccount, nWorkers)
   144  	}
   145  	return nil
   146  }
   147  
   148  func (t testMigration) MigrateAccount(
   149  	ctx context.Context,
   150  	address common.Address,
   151  	accountRegisters *registers.AccountRegisters,
   152  ) error {
   153  
   154  	if t.MigrateAccountFN != nil {
   155  		return t.MigrateAccountFN(ctx, address, accountRegisters)
   156  	}
   157  	return nil
   158  }
   159  
   160  func (t testMigration) Close() error {
   161  	if t.CloseFN != nil {
   162  		return t.CloseFN()
   163  	}
   164  	return nil
   165  }