From a2acfef6216ab61d3bd86f53755b558536641930 Mon Sep 17 00:00:00 2001
From: utahta <uekusa.yuta@gmail.com>
Date: Fri, 23 Aug 2024 16:00:44 +0900
Subject: [PATCH] fix test for lsp

---
 compiler/compiler.go          | 11 +++++++++++
 lsp/server/completion_test.go | 20 ++++++++++----------
 lsp/server/definition.go      |  2 +-
 lsp/server/handler_test.go    | 10 +++++-----
 4 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/compiler/compiler.go b/compiler/compiler.go
index e2f823fd..3dc80918 100644
--- a/compiler/compiler.go
+++ b/compiler/compiler.go
@@ -28,6 +28,7 @@ import (
 type Compiler struct {
 	importPaths  []string
 	manualImport bool
+	importRule   bool
 }
 
 // Option represents compiler option.
@@ -50,6 +51,13 @@ func ManualImportOption() Option {
 	}
 }
 
+// ImportRuleOption used to reference proto files imported by grpc.federation.file.import rule.
+func ImportRuleOption() Option {
+	return func(c *Compiler) {
+		c.importRule = true
+	}
+}
+
 // New creates compiler instance.
 func New() *Compiler {
 	return &Compiler{}
@@ -183,6 +191,9 @@ func (c *Compiler) Compile(ctx context.Context, file *source.File, opts ...Optio
 	}
 	files := []string{relPath}
 	files = append(files, file.Imports()...)
+	if c.importRule {
+		files = append(files, file.ImportsByImportRule()...)
+	}
 	linkedFiles, err := compiler.Compile(ctx, files...)
 	if err != nil {
 		return nil, &CompilerError{Err: err, ErrWithPos: r.errs}
diff --git a/lsp/server/completion_test.go b/lsp/server/completion_test.go
index 49c6d06f..bf70dc7f 100644
--- a/lsp/server/completion_test.go
+++ b/lsp/server/completion_test.go
@@ -24,8 +24,8 @@ func TestCompletion(t *testing.T) {
 	completer := server.NewCompleter(compiler.New(), slog.New(slog.NewJSONHandler(os.Stdout, nil)))
 	t.Run("method", func(t *testing.T) {
 		// resolver.method value position of Post in service.proto file
-		_, candidates, err := completer.Completion(ctx, nil, path, file, source.Position{
-			Line: 39,
+		_, candidates, err := completer.Completion(ctx, []string{"testdata"}, path, file, source.Position{
+			Line: 40,
 			Col:  19,
 		})
 		if err != nil {
@@ -44,8 +44,8 @@ func TestCompletion(t *testing.T) {
 
 	t.Run("request.field", func(t *testing.T) {
 		// resolver.request.field value position of Post in service.proto file
-		_, candidates, err := completer.Completion(ctx, nil, path, file, source.Position{
-			Line: 40,
+		_, candidates, err := completer.Completion(ctx, []string{"testdata"}, path, file, source.Position{
+			Line: 41,
 			Col:  28,
 		})
 		if err != nil {
@@ -61,8 +61,8 @@ func TestCompletion(t *testing.T) {
 
 	t.Run("request.by", func(t *testing.T) {
 		// resolver.request.by value position os Post in service.proto file
-		_, candidates, err := completer.Completion(ctx, nil, path, file, source.Position{
-			Line: 40,
+		_, candidates, err := completer.Completion(ctx, []string{"testdata"}, path, file, source.Position{
+			Line: 41,
 			Col:  38,
 		})
 		if err != nil {
@@ -78,8 +78,8 @@ func TestCompletion(t *testing.T) {
 
 	t.Run("filter response", func(t *testing.T) {
 		// resolver.response.field value position of Post in service.proto file
-		_, candidates, err := completer.Completion(ctx, nil, path, file, source.Position{
-			Line: 43,
+		_, candidates, err := completer.Completion(ctx, []string{"testdata"}, path, file, source.Position{
+			Line: 44,
 			Col:  28,
 		})
 		if err != nil {
@@ -96,8 +96,8 @@ func TestCompletion(t *testing.T) {
 
 	t.Run("message", func(t *testing.T) {
 		// def[2].message value position of Post in service.proto file
-		_, candidates, err := completer.Completion(ctx, nil, path, file, source.Position{
-			Line: 47,
+		_, candidates, err := completer.Completion(ctx, []string{"testdata"}, path, file, source.Position{
+			Line: 48,
 			Col:  17,
 		})
 		if err != nil {
diff --git a/lsp/server/definition.go b/lsp/server/definition.go
index 4fde7f8c..fad4b62d 100644
--- a/lsp/server/definition.go
+++ b/lsp/server/definition.go
@@ -52,7 +52,7 @@ func (h *Handler) definitionWithLink(ctx context.Context, params *protocol.Defin
 		return nil, nil
 	}
 	h.logger.Info("node", slog.String("text", nodeInfo.RawText()))
-	protoFiles, err := h.compiler.Compile(ctx, file, compiler.ImportPathOption(h.importPaths...))
+	protoFiles, err := h.compiler.Compile(ctx, file, compiler.ImportPathOption(h.importPaths...), compiler.ImportRuleOption())
 	if err != nil {
 		return nil, err
 	}
diff --git a/lsp/server/handler_test.go b/lsp/server/handler_test.go
index 4d031fc8..b3daaba6 100644
--- a/lsp/server/handler_test.go
+++ b/lsp/server/handler_test.go
@@ -175,7 +175,7 @@ func TestHandler_Definition(t *testing.T) {
 						URI: mustTestdataAbs(t, "testdata/service.proto"),
 					},
 					Position: protocol.Position{
-						Line:      24,
+						Line:      25,
 						Character: 15,
 					},
 				},
@@ -184,8 +184,8 @@ func TestHandler_Definition(t *testing.T) {
 				{
 					URI: mustTestdataAbs(t, "testdata/service.proto"),
 					Range: protocol.Range{
-						Start: protocol.Position{Line: 32, Character: 8},
-						End:   protocol.Position{Line: 32, Character: 12},
+						Start: protocol.Position{Line: 33, Character: 8},
+						End:   protocol.Position{Line: 33, Character: 12},
 					},
 				},
 			},
@@ -198,7 +198,7 @@ func TestHandler_Definition(t *testing.T) {
 						URI: mustTestdataAbs(t, "testdata/service.proto"),
 					},
 					Position: protocol.Position{
-						Line:      38,
+						Line:      39,
 						Character: 19,
 					},
 				},
@@ -221,7 +221,7 @@ func TestHandler_Definition(t *testing.T) {
 						URI: mustTestdataAbs(t, "testdata/service.proto"),
 					},
 					Position: protocol.Position{
-						Line:      19,
+						Line:      20,
 						Character: 0,
 					},
 				},