package source import ( "bytes" "fmt" "go/ast" "go/format" "go/parser" "go/token" "testing" "golang.org/x/tools/internal/lsp/diff" ) var fset = token.NewFileSet() func parse(t *testing.T, name, in string) *ast.File { file, err := parser.ParseFile(fset, name, in, parser.ParseComments) if err != nil { t.Fatalf("%s parse: %v", name, err) } return file } func print(t *testing.T, name string, f *ast.File) string { var buf bytes.Buffer if err := format.Node(&buf, fset, f); err != nil { t.Fatalf("%s gofmt: %v", name, err) } return buf.String() } type test struct { name string renamedPkg string pkg string in string want []importInfo unchanged bool // Expect added/deleted return value to be false. } type importInfo struct { name string path string } var addTests = []test{ { name: "leave os alone", pkg: "os", in: `package main import ( "os" ) `, want: []importInfo{ importInfo{ name: "", path: "os", }, }, unchanged: true, }, { name: "package statement only", pkg: "os", in: `package main `, want: []importInfo{ importInfo{ name: "", path: "os", }, }, }, { name: "package statement no new line", pkg: "os", in: `package main`, want: []importInfo{ importInfo{ name: "", path: "os", }, }, }, { // Issue 33721: add import statement after package declaration preceded by comments. name: "issue 33721 package statement comments before", pkg: "os", in: `// Here is a comment before package main `, want: []importInfo{ importInfo{ name: "", path: "os", }, }, }, { name: "package statement comments same line", pkg: "os", in: `package main // Here is a comment after `, want: []importInfo{ importInfo{ name: "", path: "os", }, }, }, { name: "package statement comments before and after", pkg: "os", in: `// Here is a comment before package main // Here is a comment after`, want: []importInfo{ importInfo{ name: "", path: "os", }, }, }, { name: "package statement multiline comments", pkg: "os", in: `package main /* This is a multiline comment and it extends further down*/`, want: []importInfo{ importInfo{ name: "", path: "os", }, }, }, { name: "import c", pkg: "os", in: `package main import "C" `, want: []importInfo{ importInfo{ name: "", path: "os", }, importInfo{ name: "", path: "C", }, }, }, { name: "existing imports", pkg: "os", in: `package main import "io" `, want: []importInfo{ importInfo{ name: "", path: "os", }, importInfo{ name: "", path: "io", }, }, }, { name: "existing imports with comment", pkg: "os", in: `package main import "io" // A comment `, want: []importInfo{ importInfo{ name: "", path: "os", }, importInfo{ name: "", path: "io", }, }, }, { name: "existing imports multiline comment", pkg: "os", in: `package main import "io" /* A comment that extends */ `, want: []importInfo{ importInfo{ name: "", path: "os", }, importInfo{ name: "", path: "io", }, }, }, { name: "renamed import", renamedPkg: "o", pkg: "os", in: `package main `, want: []importInfo{ importInfo{ name: "o", path: "os", }, }, }, } func TestAddImport(t *testing.T) { for _, test := range addTests { file := parse(t, test.name, test.in) var before bytes.Buffer ast.Fprint(&before, fset, file, nil) edits, err := addNamedImport(fset, file, test.renamedPkg, test.pkg) if err != nil && !test.unchanged { t.Errorf("error adding import: %s", err) continue } // Apply the edits and parse the file. got := applyEdits(test.in, edits) gotFile := parse(t, test.name, got) compareImports(t, fmt.Sprintf("first run: %s:\n", test.name), gotFile.Imports, test.want) // AddNamedImport should be idempotent. Verify that by calling it again, // expecting no change to the AST, and the returned added value to always be false. edits, err = addNamedImport(fset, gotFile, test.renamedPkg, test.pkg) if err != nil && !test.unchanged { t.Errorf("error adding import: %s", err) continue } // Apply the edits and parse the file. got = applyEdits(got, edits) gotFile = parse(t, test.name, got) compareImports(t, test.name, gotFile.Imports, test.want) } } func TestDoubleAddNamedImport(t *testing.T) { name := "doublenamedimport" in := "package main\n" file := parse(t, name, in) // Add a named import edits, err := addNamedImport(fset, file, "o", "os") if err != nil { t.Errorf("error adding import: %s", err) return } got := applyEdits(in, edits) gotFile := parse(t, name, got) // Add a second named import edits, err = addNamedImport(fset, gotFile, "i", "io") if err != nil { t.Errorf("error adding import: %s", err) return } got = applyEdits(got, edits) gotFile = parse(t, name, got) want := []importInfo{ importInfo{ name: "o", path: "os", }, importInfo{ name: "i", path: "io", }, } compareImports(t, "", gotFile.Imports, want) } func compareImports(t *testing.T, prefix string, got []*ast.ImportSpec, want []importInfo) { if len(got) != len(want) { t.Errorf("%s\ngot %d imports\nwant %d", prefix, len(got), len(want)) return } for _, imp := range got { name := importName(imp) path := importPath(imp) found := false for _, want := range want { if want.name == name && want.path == path { found = true break } } if !found { t.Errorf("%s\n\ngot unexpected import: name: %q,path: %q", prefix, name, path) continue } } } func applyEdits(contents string, edits []diff.TextEdit) string { res := contents // Apply the edits from the end of the file forward // to preserve the offsets for i := len(edits) - 1; i >= 0; i-- { edit := edits[i] start := edit.Span.Start().Offset() end := edit.Span.End().Offset() tmp := res[0:start] + edit.NewText res = tmp + res[end:] } return res }