Fix import location

kitagry · Apr 29, 2022 · 5efa205 · 5efa205
1 parent 458eeea
commit 5efa205
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 1 deletion.
diff --git a/langserver/internal/source/completion.go b/langserver/internal/source/completion.go
@@ -380,13 +380,21 @@ func (p *Project) listImportCompletionItems(location *ast.Location) []Completion
 		}
 	}
 
+	locationForImport := &ast.Location{
+		Text:   location.Text,
+		File:   location.File,
+		Row:    location.Row,
+		Col:    1,
+		Offset: location.Offset,
+	}
+
 	result := make([]CompletionItem, 0, len(refs))
 	for _, r := range refs {
 		if !inRef(r, alreadyExistPackages) {
 			result = append(result, CompletionItem{
 				Label:    fmt.Sprintf("import %s", r.String()),
 				Kind:     ImportItem,
-				TextEdit: createTextEdit(location, fmt.Sprintf("import %s", r.String())),
+				TextEdit: createTextEdit(locationForImport, fmt.Sprintf("import %s", r.String())),
 			})
 		}
 	}

diff --git a/langserver/internal/source/completion_test.go b/langserver/internal/source/completion_test.go
@@ -10,6 +10,7 @@ import (
 
 type completionTestCase struct {
 	files          map[string]source.File
+	updateFile     map[string]source.File
 	createLocation createLocationFunc
 	expectItems    []source.CompletionItem
 }
@@ -40,6 +41,39 @@ func TestProject_ListCompletionItemsStrict(t *testing.T) {
 				},
 			},
 		},
+		"Should list import library location is 1": {
+			files: map[string]source.File{
+				"src.rego": {
+					RawText: `package src
+
+
+`,
+				},
+				"lib.rego": {
+					RawText: `package lib`,
+				},
+			},
+			updateFile: map[string]source.File{
+				"src.rego": {
+					RawText: `package src
+
+im
+`,
+				},
+			},
+			createLocation: createLocation(3, 2, "src.rego"),
+			expectItems: []source.CompletionItem{
+				{
+					Label: "import data.lib",
+					Kind:  source.ImportItem,
+					TextEdit: &source.TextEdit{
+						Row:  3,
+						Col:  1,
+						Text: "import data.lib",
+					},
+				},
+			},
+		},
 		"Should not list already imported library": {
 			files: map[string]source.File{
 				"src.rego": {
@@ -184,6 +218,13 @@ violation[msg] {
 				t.Fatal(err)
 			}
 
+			for path, file := range tt.updateFile {
+				err := project.UpdateFile(path, file.RawText, file.Version)
+				if err != nil {
+					t.Fatal(err)
+				}
+			}
+
 			location := tt.createLocation(tt.files)
 			got, err := project.ListCompletionItems(location)
 			if err != nil {
@@ -508,6 +549,13 @@ default is_test = true`,
 						t.Fatal(err)
 					}
 
+					for path, file := range tt.updateFile {
+						err := project.UpdateFile(path, file.RawText, file.Version)
+						if err != nil {
+							t.Fatal(err)
+						}
+					}
+
 					location := tt.createLocation(tt.files)
 					got, err := project.ListCompletionItems(location)
 					if err != nil {