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

     1  package migrations
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"strings"
     7  	"testing"
     8  
     9  	"github.com/onflow/cadence/runtime/common"
    10  	"github.com/rs/zerolog"
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/require"
    13  
    14  	"github.com/onflow/flow-go/cmd/util/ledger/util/registers"
    15  	"github.com/onflow/flow-go/model/flow"
    16  )
    17  
    18  const contractA = `
    19  access(all) contract A {
    20      access(all) fun foo() {}
    21  }`
    22  
    23  const updatedContractA = `
    24  access(all) contract A {
    25      access(all) fun bar() {}
    26  }`
    27  
    28  const contractB = `
    29  access(all) contract B {
    30      access(all) fun foo() {}
    31  }`
    32  
    33  const updatedContractB = `
    34  access(all) contract B {
    35      access(all) fun bar() {}
    36  }`
    37  
    38  type errorLogWriter struct {
    39  	logs []string
    40  }
    41  
    42  var _ io.Writer = &errorLogWriter{}
    43  
    44  const errorLogPrefix = "{\"level\":\"error\""
    45  
    46  func (l *errorLogWriter) Write(bytes []byte) (int, error) {
    47  	logStr := string(bytes)
    48  
    49  	// Ignore non-error logs
    50  	if !strings.HasPrefix(logStr, errorLogPrefix) {
    51  		return 0, nil
    52  	}
    53  
    54  	l.logs = append(l.logs, logStr)
    55  	return len(bytes), nil
    56  }
    57  
    58  func TestChangeContractCodeMigration(t *testing.T) {
    59  	t.Parallel()
    60  
    61  	const chainID = flow.Emulator
    62  	addressGenerator := chainID.Chain().NewAddressGenerator()
    63  
    64  	address1, err := addressGenerator.NextAddress()
    65  	require.NoError(t, err)
    66  
    67  	address2, err := addressGenerator.NextAddress()
    68  	require.NoError(t, err)
    69  
    70  	ctx := context.Background()
    71  
    72  	t.Run("no contracts", func(t *testing.T) {
    73  		t.Parallel()
    74  
    75  		writer := &errorLogWriter{}
    76  		log := zerolog.New(writer)
    77  
    78  		rwf := &testReportWriterFactory{}
    79  
    80  		options := StagedContractsMigrationOptions{
    81  			ChainID:            flow.Emulator,
    82  			VerboseErrorOutput: true,
    83  		}
    84  		migration := NewStagedContractsMigration("test", "test", log, rwf, options)
    85  
    86  		registersByAccount := registers.NewByAccount()
    87  
    88  		err := migration.InitMigration(log, registersByAccount, 1)
    89  		require.NoError(t, err)
    90  
    91  		err = migration.MigrateAccount(
    92  			ctx,
    93  			common.Address(address1),
    94  			registersByAccount.AccountRegisters(string(address1[:])),
    95  		)
    96  		require.NoError(t, err)
    97  
    98  		err = migration.Close()
    99  		require.NoError(t, err)
   100  
   101  		require.Empty(t, writer.logs)
   102  	})
   103  
   104  	t.Run("1 contract - dont migrate", func(t *testing.T) {
   105  		t.Parallel()
   106  
   107  		writer := &errorLogWriter{}
   108  		log := zerolog.New(writer)
   109  
   110  		rwf := &testReportWriterFactory{}
   111  
   112  		options := StagedContractsMigrationOptions{
   113  			ChainID:            flow.Emulator,
   114  			VerboseErrorOutput: true,
   115  		}
   116  		migration := NewStagedContractsMigration("test", "test", log, rwf, options)
   117  
   118  		registersByAccount, err := registersForStagedContracts(
   119  			StagedContract{
   120  				Address: common.Address(address1),
   121  				Contract: Contract{
   122  					Name: "A",
   123  					Code: []byte(contractA),
   124  				},
   125  			},
   126  		)
   127  		require.NoError(t, err)
   128  
   129  		err = migration.InitMigration(log, registersByAccount, 1)
   130  		require.NoError(t, err)
   131  
   132  		owner1 := string(address1[:])
   133  
   134  		accountRegisters1 := registersByAccount.AccountRegisters(owner1)
   135  
   136  		err = migration.MigrateAccount(
   137  			ctx,
   138  			common.Address(address1),
   139  			accountRegisters1,
   140  		)
   141  		require.NoError(t, err)
   142  
   143  		require.Equal(t, 1, registersByAccount.AccountCount())
   144  		require.Equal(t, 1, accountRegisters1.Count())
   145  		require.Equal(t,
   146  			contractA,
   147  			contractCode(t, registersByAccount, owner1, "A"),
   148  		)
   149  
   150  		err = migration.Close()
   151  		require.NoError(t, err)
   152  
   153  		require.Empty(t, writer.logs)
   154  	})
   155  
   156  	t.Run("1 contract - migrate", func(t *testing.T) {
   157  		t.Parallel()
   158  
   159  		writer := &errorLogWriter{}
   160  		log := zerolog.New(writer)
   161  
   162  		rwf := &testReportWriterFactory{}
   163  
   164  		options := StagedContractsMigrationOptions{
   165  			ChainID:            flow.Emulator,
   166  			VerboseErrorOutput: true,
   167  		}
   168  		migration := NewStagedContractsMigration("test", "test", log, rwf, options).
   169  			WithStagedContractUpdates([]StagedContract{
   170  				{
   171  					Address: common.Address(address1),
   172  					Contract: Contract{
   173  						Name: "A",
   174  						Code: []byte(updatedContractA),
   175  					},
   176  				},
   177  			})
   178  
   179  		registersByAccount, err := registersForStagedContracts(
   180  			StagedContract{
   181  				Address: common.Address(address1),
   182  				Contract: Contract{
   183  					Name: "A",
   184  					Code: []byte(contractA),
   185  				},
   186  			},
   187  		)
   188  		require.NoError(t, err)
   189  
   190  		err = migration.InitMigration(log, registersByAccount, 1)
   191  		require.NoError(t, err)
   192  
   193  		owner1 := string(address1[:])
   194  
   195  		accountRegisters1 := registersByAccount.AccountRegisters(owner1)
   196  
   197  		err = migration.MigrateAccount(
   198  			ctx,
   199  			common.Address(address1),
   200  			accountRegisters1,
   201  		)
   202  		require.NoError(t, err)
   203  
   204  		require.Equal(t, 1, registersByAccount.AccountCount())
   205  		require.Equal(t, 1, accountRegisters1.Count())
   206  		require.Equal(t,
   207  			updatedContractA,
   208  			contractCode(t, registersByAccount, owner1, "A"),
   209  		)
   210  
   211  		err = migration.Close()
   212  		require.NoError(t, err)
   213  
   214  		require.Empty(t, writer.logs)
   215  	})
   216  
   217  	t.Run("2 contracts - migrate 1", func(t *testing.T) {
   218  		t.Parallel()
   219  
   220  		writer := &errorLogWriter{}
   221  		log := zerolog.New(writer)
   222  
   223  		rwf := &testReportWriterFactory{}
   224  
   225  		options := StagedContractsMigrationOptions{
   226  			ChainID:            flow.Emulator,
   227  			VerboseErrorOutput: true,
   228  		}
   229  		migration := NewStagedContractsMigration("test", "test", log, rwf, options).
   230  			WithStagedContractUpdates([]StagedContract{
   231  				{
   232  					Address: common.Address(address1),
   233  					Contract: Contract{
   234  						Name: "A",
   235  						Code: []byte(updatedContractA),
   236  					},
   237  				},
   238  			})
   239  
   240  		registersByAccount, err := registersForStagedContracts(
   241  			StagedContract{
   242  				Address: common.Address(address1),
   243  				Contract: Contract{
   244  					Name: "A",
   245  					Code: []byte(contractA),
   246  				},
   247  			},
   248  			StagedContract{
   249  				Address: common.Address(address1),
   250  				Contract: Contract{
   251  					Name: "B",
   252  					Code: []byte(contractB),
   253  				},
   254  			},
   255  		)
   256  		require.NoError(t, err)
   257  
   258  		err = migration.InitMigration(log, registersByAccount, 1)
   259  		require.NoError(t, err)
   260  
   261  		owner1 := string(address1[:])
   262  
   263  		accountRegisters1 := registersByAccount.AccountRegisters(owner1)
   264  
   265  		err = migration.MigrateAccount(
   266  			ctx,
   267  			common.Address(address1),
   268  			accountRegisters1,
   269  		)
   270  		require.NoError(t, err)
   271  
   272  		require.Equal(t, 1, registersByAccount.AccountCount())
   273  		require.Equal(t, 2, accountRegisters1.Count())
   274  		require.Equal(t,
   275  			updatedContractA,
   276  			contractCode(t, registersByAccount, owner1, "A"),
   277  		)
   278  		require.Equal(t,
   279  			contractB,
   280  			contractCode(t, registersByAccount, owner1, "B"),
   281  		)
   282  
   283  		err = migration.Close()
   284  		require.NoError(t, err)
   285  
   286  		require.Empty(t, writer.logs)
   287  	})
   288  
   289  	t.Run("2 contracts - migrate 2", func(t *testing.T) {
   290  		t.Parallel()
   291  
   292  		writer := &errorLogWriter{}
   293  		log := zerolog.New(writer)
   294  
   295  		rwf := &testReportWriterFactory{}
   296  
   297  		options := StagedContractsMigrationOptions{
   298  			ChainID:            flow.Emulator,
   299  			VerboseErrorOutput: true,
   300  		}
   301  		migration := NewStagedContractsMigration("test", "test", log, rwf, options).
   302  			WithStagedContractUpdates([]StagedContract{
   303  				{
   304  					Address: common.Address(address1),
   305  					Contract: Contract{
   306  						Name: "A",
   307  						Code: []byte(updatedContractA),
   308  					},
   309  				},
   310  				{
   311  					Address: common.Address(address1),
   312  					Contract: Contract{
   313  						Name: "B",
   314  						Code: []byte(updatedContractB),
   315  					},
   316  				},
   317  			})
   318  
   319  		registersByAccount, err := registersForStagedContracts(
   320  			StagedContract{
   321  				Address: common.Address(address1),
   322  				Contract: Contract{
   323  					Name: "A",
   324  					Code: []byte(contractA),
   325  				},
   326  			},
   327  			StagedContract{
   328  				Address: common.Address(address1),
   329  				Contract: Contract{
   330  					Name: "B",
   331  					Code: []byte(contractB),
   332  				},
   333  			},
   334  		)
   335  		require.NoError(t, err)
   336  
   337  		err = migration.InitMigration(log, registersByAccount, 1)
   338  		require.NoError(t, err)
   339  
   340  		owner1 := string(address1[:])
   341  
   342  		accountRegisters1 := registersByAccount.AccountRegisters(owner1)
   343  
   344  		err = migration.MigrateAccount(
   345  			ctx,
   346  			common.Address(address1),
   347  			accountRegisters1,
   348  		)
   349  		require.NoError(t, err)
   350  
   351  		require.Equal(t, 1, registersByAccount.AccountCount())
   352  		require.Equal(t, 2, accountRegisters1.Count())
   353  		require.Equal(t,
   354  			updatedContractA,
   355  			contractCode(t, registersByAccount, owner1, "A"),
   356  		)
   357  		require.Equal(t,
   358  			updatedContractB,
   359  			contractCode(t, registersByAccount, owner1, "B"),
   360  		)
   361  
   362  		err = migration.Close()
   363  		require.NoError(t, err)
   364  
   365  		require.Empty(t, writer.logs)
   366  	})
   367  
   368  	t.Run("2 contracts on different accounts - migrate 1", func(t *testing.T) {
   369  		t.Parallel()
   370  
   371  		writer := &errorLogWriter{}
   372  		log := zerolog.New(writer)
   373  
   374  		rwf := &testReportWriterFactory{}
   375  
   376  		options := StagedContractsMigrationOptions{
   377  			ChainID:            flow.Emulator,
   378  			VerboseErrorOutput: true,
   379  		}
   380  		migration := NewStagedContractsMigration("test", "test", log, rwf, options).
   381  			WithStagedContractUpdates([]StagedContract{
   382  				{
   383  					Address: common.Address(address1),
   384  					Contract: Contract{
   385  						Name: "A",
   386  						Code: []byte(updatedContractA),
   387  					},
   388  				},
   389  			})
   390  
   391  		registersByAccount, err := registersForStagedContracts(
   392  			StagedContract{
   393  				Address: common.Address(address1),
   394  				Contract: Contract{
   395  					Name: "A",
   396  					Code: []byte(contractA),
   397  				},
   398  			},
   399  			StagedContract{
   400  				Address: common.Address(address2),
   401  				Contract: Contract{
   402  					Name: "A",
   403  					Code: []byte(contractA),
   404  				},
   405  			},
   406  		)
   407  		require.NoError(t, err)
   408  
   409  		err = migration.InitMigration(log, registersByAccount, 1)
   410  		require.NoError(t, err)
   411  
   412  		owner1 := string(address1[:])
   413  		owner2 := string(address2[:])
   414  
   415  		accountRegisters1 := registersByAccount.AccountRegisters(owner1)
   416  		accountRegisters2 := registersByAccount.AccountRegisters(owner2)
   417  
   418  		err = migration.MigrateAccount(
   419  			ctx,
   420  			common.Address(address1),
   421  			accountRegisters1,
   422  		)
   423  		require.NoError(t, err)
   424  
   425  		require.Equal(t, 2, registersByAccount.AccountCount())
   426  		require.Equal(t, 1, accountRegisters1.Count())
   427  		require.Equal(t, 1, accountRegisters2.Count())
   428  		require.Equal(t,
   429  			updatedContractA,
   430  			contractCode(t, registersByAccount, owner1, "A"),
   431  		)
   432  		require.Equal(t,
   433  			contractA,
   434  			contractCode(t, registersByAccount, owner2, "A"),
   435  		)
   436  
   437  		err = migration.Close()
   438  		require.NoError(t, err)
   439  
   440  		require.Empty(t, writer.logs)
   441  	})
   442  
   443  	t.Run("not all contracts on one account migrated", func(t *testing.T) {
   444  		t.Parallel()
   445  
   446  		writer := &errorLogWriter{}
   447  		log := zerolog.New(writer)
   448  
   449  		rwf := &testReportWriterFactory{}
   450  
   451  		options := StagedContractsMigrationOptions{
   452  			ChainID:            flow.Emulator,
   453  			VerboseErrorOutput: true,
   454  		}
   455  		migration := NewStagedContractsMigration("test", "test", log, rwf, options).
   456  			WithStagedContractUpdates([]StagedContract{
   457  				{
   458  					Address: common.Address(address1),
   459  					Contract: Contract{
   460  						Name: "A",
   461  						Code: []byte(updatedContractA),
   462  					},
   463  				},
   464  				{
   465  					Address: common.Address(address1),
   466  					Contract: Contract{
   467  						Name: "B",
   468  						Code: []byte(updatedContractB),
   469  					},
   470  				},
   471  			})
   472  
   473  		registersByAccount, err := registersForStagedContracts(
   474  			StagedContract{
   475  				Address: common.Address(address1),
   476  				Contract: Contract{
   477  					Name: "A",
   478  					Code: []byte(contractA),
   479  				},
   480  			},
   481  		)
   482  		require.NoError(t, err)
   483  
   484  		err = migration.InitMigration(log, registersByAccount, 1)
   485  		require.NoError(t, err)
   486  
   487  		owner1 := string(address1[:])
   488  
   489  		accountRegisters1 := registersByAccount.AccountRegisters(owner1)
   490  
   491  		err = migration.MigrateAccount(
   492  			ctx,
   493  			common.Address(address1),
   494  			accountRegisters1,
   495  		)
   496  		require.NoError(t, err)
   497  
   498  		require.Len(t, writer.logs, 1)
   499  		assert.Contains(t,
   500  			writer.logs[0],
   501  			`missing old code`,
   502  		)
   503  	})
   504  
   505  	t.Run("not all accounts migrated", func(t *testing.T) {
   506  		t.Parallel()
   507  
   508  		writer := &errorLogWriter{}
   509  		log := zerolog.New(writer)
   510  
   511  		rwf := &testReportWriterFactory{}
   512  
   513  		options := StagedContractsMigrationOptions{
   514  			ChainID:            flow.Emulator,
   515  			VerboseErrorOutput: true,
   516  		}
   517  		migration := NewStagedContractsMigration("test", "test", log, rwf, options).
   518  			WithStagedContractUpdates([]StagedContract{
   519  				{
   520  					Address: common.Address(address2),
   521  					Contract: Contract{
   522  						Name: "A",
   523  						Code: []byte(updatedContractA),
   524  					},
   525  				},
   526  			})
   527  
   528  		registersByAccount, err := registersForStagedContracts(
   529  			StagedContract{
   530  				Address: common.Address(address1),
   531  				Contract: Contract{
   532  					Name: "A",
   533  					Code: []byte(contractA),
   534  				},
   535  			},
   536  		)
   537  		require.NoError(t, err)
   538  
   539  		err = migration.InitMigration(log, registersByAccount, 1)
   540  		require.NoError(t, err)
   541  
   542  		owner1 := string(address1[:])
   543  
   544  		accountRegisters1 := registersByAccount.AccountRegisters(owner1)
   545  
   546  		err = migration.MigrateAccount(
   547  			ctx,
   548  			common.Address(address1),
   549  			accountRegisters1,
   550  		)
   551  		require.NoError(t, err)
   552  
   553  		err = migration.Close()
   554  		require.NoError(t, err)
   555  
   556  		require.Len(t, writer.logs, 1)
   557  		assert.Contains(t,
   558  			writer.logs[0],
   559  			`"failed to find all contract registers that need to be changed"`,
   560  		)
   561  	})
   562  }