From 97c69f7efa4ee904350c45e6c00101c92fd69f58 Mon Sep 17 00:00:00 2001
From: Jordan Ribbink <jribbink@telus.net>
Date: Wed, 17 Apr 2024 14:42:18 -0700
Subject: [PATCH] add nil check & test

---
 internal/migrate/staging_validator.go      | 10 ++--
 internal/migrate/staging_validator_test.go | 61 ++++++++++++++++++++++
 2 files changed, 66 insertions(+), 5 deletions(-)

diff --git a/internal/migrate/staging_validator.go b/internal/migrate/staging_validator.go
index 82c81b387..ba4d1c9e6 100644
--- a/internal/migrate/staging_validator.go
+++ b/internal/migrate/staging_validator.go
@@ -391,6 +391,10 @@ func (v *stagingValidator) resolveLocation(
 }
 
 func (v *stagingValidator) resolveAccountAccess(checker *sema.Checker, memberLocation common.Location) bool {
+	if checker == nil {
+		return false
+	}
+
 	checkerLocation, ok := checker.Location.(common.StringLocation)
 	if !ok {
 		return false
@@ -403,11 +407,7 @@ func (v *stagingValidator) resolveAccountAccess(checker *sema.Checker, memberLoc
 
 	// If the source code of the update is being checked, we should check account access based on the
 	// targeted network location of the contract & not the source code location
-	if checkerLocation == v.sourceCodeLocation && memberAddressLocation.Address == v.targetLocation.Address {
-		return true
-	}
-
-	return false
+	return checkerLocation == v.sourceCodeLocation && memberAddressLocation.Address == v.targetLocation.Address
 }
 
 func (v *stagingValidator) resolveAddressContractNames(address common.Address) ([]string, error) {
diff --git a/internal/migrate/staging_validator_test.go b/internal/migrate/staging_validator_test.go
index 48db6e70f..fd2e5070b 100644
--- a/internal/migrate/staging_validator_test.go
+++ b/internal/migrate/staging_validator_test.go
@@ -288,4 +288,65 @@ func Test_StagingValidator(t *testing.T) {
 		err := validator.ValidateContractUpdate(location, sourceCodeLocation, []byte(newContract))
 		require.NoError(t, err)
 	})
+
+	t.Run("resolves account access correctly", func(t *testing.T) {
+		location := common.NewAddressLocation(nil, common.Address{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01}, "Test")
+		sourceCodeLocation := common.StringLocation("./Test.cdc")
+		oldContract := `
+		import ImpContract from 0x01
+		pub contract Test {
+			pub fun test() {}
+		}`
+		newContract := `
+		import ImpContract from 0x01
+		access(all) contract Test {
+			access(all) fun test() {}
+			init() {
+				ImpContract.test()
+			}
+		}`
+		impContract := `
+		access(all) contract ImpContract {
+			access(account) fun test() {}
+			init() {}
+		}`
+		mockScriptResultString, err := cadence.NewString(impContract)
+		require.NoError(t, err)
+
+		mockAccount := &flow.Account{
+			Address: flow.HexToAddress("01"),
+			Balance: 1000,
+			Keys:    nil,
+			Contracts: map[string][]byte{
+				"Test": []byte(oldContract),
+			},
+		}
+
+		// setup mocks
+		require.NoError(t, rw.WriteFile(sourceCodeLocation.String(), []byte(newContract), 0o644))
+		srv.GetAccount.Run(func(args mock.Arguments) {
+			require.Equal(t, flow.HexToAddress("01"), args.Get(1).(flow.Address))
+		}).Return(mockAccount, nil)
+		srv.Network.Return(config.Network{
+			Name: "testnet",
+		}, nil)
+		srv.ExecuteScript.Run(func(args mock.Arguments) {
+			script := args.Get(1).(flowkit.Script)
+
+			assert.Equal(t, templates.GenerateGetStagedContractCodeScript(MigrationContractStagingAddress("testnet")), script.Code)
+
+			assert.Equal(t, 2, len(script.Args))
+			actualContractAddressArg, actualContractNameArg := script.Args[0], script.Args[1]
+
+			contractName, _ := cadence.NewString("ImpContract")
+			contractAddr := cadence.NewAddress(flow.HexToAddress("01"))
+			assert.Equal(t, contractName, actualContractNameArg)
+			assert.Equal(t, contractAddr, actualContractAddressArg)
+		}).Return(cadence.NewOptional(mockScriptResultString), nil)
+
+		// validate
+		validator := newStagingValidator(srv.Mock, state)
+		err = validator.ValidateContractUpdate(location, sourceCodeLocation, []byte(newContract))
+		require.NoError(t, err)
+	})
 }