update package to replace both formats to get tests passing

onflow · Feb 14, 2025 · 613d745 · 613d745
1 parent 6e1fc44
commit 613d745
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 34 deletions.
diff --git a/lib/go/contracts/contracts.go b/lib/go/contracts/contracts.go
@@ -193,12 +193,11 @@ func FlowServiceAccount(env templates.Environment) []byte {
 	code := assets.MustAssetString(flowServiceAccountFilename)
 
 	if env.FlowExecutionParametersAddress == "" {
-
 		// Remove the import of FlowExecutionParameters
 		code = strings.ReplaceAll(
 			code,
-			"import FlowExecutionParameters from \"FlowExecutionParameters\"",
-			"//import FlowExecutionParameters from \"FlowExecutionParameters\"",
+			"import \"FlowExecutionParameters\"",
+			"//import \"FlowExecutionParameters\"",
 		)
 
 		// Replace the metering getter functions
@@ -328,16 +327,20 @@ func TESTFlowStakingCollection(
 ) []byte {
 	code := assets.MustAssetString(flowStakingCollectionFilename)
 
-	code = strings.ReplaceAll(code, placeholderFungibleTokenAddress, withHexPrefix(fungibleTokenAddress))
-	code = strings.ReplaceAll(code, placeholderBurnerAddress, withHexPrefix(storageFeesAddress))
-	code = strings.ReplaceAll(code, placeholderFlowTokenAddress, withHexPrefix(flowTokenAddress))
-	code = strings.ReplaceAll(code, placeholderIDTableAddress, withHexPrefix(idTableAddress))
-	code = strings.ReplaceAll(code, placeholderStakingProxyAddress, withHexPrefix(stakingProxyAddress))
-	code = strings.ReplaceAll(code, placeholderLockedTokensAddress, withHexPrefix(lockedTokensAddress))
-	code = strings.ReplaceAll(code, placeholderStorageFeesAddress, withHexPrefix(storageFeesAddress))
-	code = strings.ReplaceAll(code, placeholderQCAddr, withHexPrefix(qcAddress))
-	code = strings.ReplaceAll(code, placeholderDKGAddr, withHexPrefix(dkgAddress))
-	code = strings.ReplaceAll(code, placeholderEpochAddr, withHexPrefix(epochAddress))
+	env := templates.Environment{
+		FungibleTokenAddress:     fungibleTokenAddress,
+		FlowTokenAddress:         flowTokenAddress,
+		IDTableAddress:           idTableAddress,
+		StakingProxyAddress:      stakingProxyAddress,
+		LockedTokensAddress:      lockedTokensAddress,
+		StorageFeesAddress:       storageFeesAddress,
+		QuorumCertificateAddress: qcAddress,
+		DkgAddress:               dkgAddress,
+		EpochAddress:             epochAddress,
+		BurnerAddress:            storageFeesAddress,
+	}
+
+	code = templates.ReplaceAddresses(code, env)
 
 	code = strings.ReplaceAll(code, "access(self) fun getTokens", "access(all) fun getTokens")
 	code = strings.ReplaceAll(code, "access(self) fun depositTokens", "access(all) fun depositTokens")
@@ -348,23 +351,13 @@ func TESTFlowStakingCollection(
 func TestFlowFees(fungibleTokenAddress, flowTokenAddress, storageFeesAddress string) []byte {
 	code := assets.MustAssetString(flowFeesFilename)
 
-	code = strings.ReplaceAll(
-		code,
-		placeholderFungibleTokenAddress,
-		withHexPrefix(fungibleTokenAddress),
-	)
-
-	code = strings.ReplaceAll(
-		code,
-		placeholderFlowTokenAddress,
-		withHexPrefix(flowTokenAddress),
-	)
+	env := templates.Environment{
+		FungibleTokenAddress: fungibleTokenAddress,
+		FlowTokenAddress:     flowTokenAddress,
+		StorageFeesAddress:   storageFeesAddress,
+	}
 
-	code = strings.ReplaceAll(
-		code,
-		placeholderStorageFeesAddress,
-		withHexPrefix(storageFeesAddress),
-	)
+	code = templates.ReplaceAddresses(code, env)
 
 	return []byte(code)
 }
diff --git a/lib/go/templates/templates.go b/lib/go/templates/templates.go
@@ -81,11 +81,19 @@ func ReplaceAddress(code, placeholder, replacement string) string {
 	placeholderWithoutQuotes := placeholder[1 : len(placeholder)-1]
 
 	if len(replacement) > 0 {
-		code = strings.ReplaceAll(
-			code,
-			placeholder,
-			placeholderWithoutQuotes+" from "+withHexPrefix(replacement),
-		)
+		if strings.Contains(code, placeholderWithoutQuotes+" from "+placeholder) {
+			code = strings.ReplaceAll(
+				code,
+				placeholder,
+				withHexPrefix(replacement),
+			)
+		} else {
+			code = strings.ReplaceAll(
+				code,
+				placeholder,
+				placeholderWithoutQuotes+" from "+withHexPrefix(replacement),
+			)
+		}
 	}
 	return code
 }