Reuse the replaceImports method of EmulatorBackend

onflow · Sep 29, 2023 · 644e5ac · 644e5ac
1 parent 961ca96
commit 644e5ac
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 54 deletions.
diff --git a/test/emulator_backend.go b/test/emulator_backend.go
@@ -675,15 +675,16 @@ func (e *EmulatorBackend) replaceImports(code string) string {
 		}
 
 		var address common.Address
+		var found bool
 		if len(importDeclaration.Identifiers) > 0 {
-			address, ok = e.contracts[importDeclaration.Identifiers[0].Identifier]
-			if !ok {
-				// keep import statement it as-is
-				sb.WriteString(code[prevImportDeclEnd:importDeclEnd])
-				continue
-			}
+			address, found = e.contracts[importDeclaration.Identifiers[0].Identifier]
 		} else {
-			address = e.contracts[location.String()]
+			address, found = e.contracts[location.String()]
+		}
+		if !found {
+			// keep import statement it as-is
+			sb.WriteString(code[prevImportDeclEnd:importDeclEnd])
+			continue
 		}
 
 		var importStr string

diff --git a/test/test_runner.go b/test/test_runner.go
@@ -287,52 +287,8 @@ func (r *TestRunner) RunTests(script string) (results Results, err error) {
 	return results, err
 }
 
-func (r *TestRunner) replaceImports(program *ast.Program, code string) string {
-	sb := strings.Builder{}
-	importDeclEnd := 0
-
-	for _, importDeclaration := range program.ImportDeclarations() {
-		prevImportDeclEnd := importDeclEnd
-		importDeclEnd = importDeclaration.EndPos.Offset + 1
-
-		location, ok := importDeclaration.Location.(common.StringLocation)
-		if !ok {
-			// keep the import statement it as-is
-			sb.WriteString(code[prevImportDeclEnd:importDeclEnd])
-			continue
-		}
-
-		var address common.Address
-		var found bool
-		if len(importDeclaration.Identifiers) > 0 {
-			address, found = r.contracts[importDeclaration.Identifiers[0].Identifier]
-		} else {
-			address, found = r.contracts[location.String()]
-		}
-		if !found {
-			// keep import statement it as-is
-			sb.WriteString(code[prevImportDeclEnd:importDeclEnd])
-			continue
-		}
-
-		var importStr string
-		if strings.Contains(importDeclaration.String(), "from") {
-			importStr = fmt.Sprintf("0x%s", address)
-		} else {
-			// Imports of the form `import "FungibleToken"` should be
-			// expanded to `import FungibleToken from 0xee82856bf20e2aa6`
-			importStr = fmt.Sprintf("%s from 0x%s", location, address)
-		}
-
-		locationStart := importDeclaration.LocationPos.Offset
-
-		sb.WriteString(code[prevImportDeclEnd:locationStart])
-		sb.WriteString(importStr)
-	}
-
-	sb.WriteString(code[importDeclEnd:])
-
-	return sb.String()
+func (r *TestRunner) replaceImports(code string) string {
+	return r.backend.replaceImports(code)
 }
 
 func (r *TestRunner) runTestSetup(inter *interpreter.Interpreter) error {
@@ -467,7 +423,6 @@ func (r *TestRunner) parseCheckAndInterpret(script string) (*interpreter.Program
 	if err != nil {
 		panic(err)
 	}
-	script = r.replaceImports(code, script)
 
 	for _, funcDecl := range code.FunctionDeclarations() {
 		funcName := funcDecl.Identifier.Identifier
@@ -485,6 +440,7 @@ func (r *TestRunner) parseCheckAndInterpret(script string) (*interpreter.Program
 		}
 	}
 
+	script = r.replaceImports(script)
 	program, err := r.testRuntime.ParseAndCheckProgram([]byte(script), ctx)
 	if err != nil {
 		return nil, nil, err