fix template import aliasing

lib/go/templates/templates.go

@@ -1,7 +1,6 @@
 package templates
 
 import (
-	"fmt"
 	"regexp"
 
 	"github.com/onflow/flow-go-sdk"
@@ -13,19 +12,12 @@ var (
 	placeholderMigrationContractStaging = regexp.MustCompile(`"MigrationContractStaging"`)
 )
 
+// Replaces the import alias of the form `import "MigrationContractStaging"` in the given code with the given address,
+// resuting in a static import statement of the form `import MigrationContractStaging from 0xADDRESS`.
 func replaceMigrationContractStagingImports(code string, migrationContractStagingAddress flow.Address) []byte {
-	code = placeholderMigrationContractStaging.ReplaceAllString(code, "0x"+migrationContractStagingAddress.String())
+	code = placeholderMigrationContractStaging.ReplaceAllString(
+		code,
+		"MigrationContractStaging from 0x"+migrationContractStagingAddress.String(),
+	)
 	return []byte(code)
 }
-
-func withHexPrefix(address string) string {
-	if address == "" {
-		return ""
-	}
-
-	if address[0:2] == "0x" {
-		return address
-	}
-
-	return fmt.Sprintf("0x%s", address)
-}

lib/go/templates/templates_test.go

@@ -1,18 +1,34 @@
 package templates_test
 
 import (
+	"strings"
 	"testing"
 
 	"github.com/onflow/contract-updater/lib/go/templates"
 	"github.com/onflow/flow-go-sdk/test"
 	"github.com/stretchr/testify/assert"
 )
 
+func TestIsValidatedScript(t *testing.T) {
+	addresses := test.AddressGenerator()
+	contractAlias := addresses.New()
+
+	template := templates.GenerateIsValidatedScript(contractAlias)
+	assert.NotNil(t, template)
+
+	importLine := strings.Split(string(template), "\n")[0]
+	expectedImportLine := `import MigrationContractStaging from 0x` + contractAlias.String()
+	assert.Equal(t, expectedImportLine, importLine)
+}
+
 func TestGenerateStageContractScript(t *testing.T) {
 	addresses := test.AddressGenerator()
 	contractAlias := addresses.New()
 
 	template := templates.GenerateStageContractScript(contractAlias)
 	assert.NotNil(t, template)
-	assert.Contains(t, string(template), contractAlias.String())
+
+	importLine := strings.Split(string(template), "\n")[0]
+	expectedImportLine := `import MigrationContractStaging from 0x` + contractAlias.String()
+	assert.Equal(t, expectedImportLine, importLine)
 }