From 87018ac365d7196fb81fdb5b468c2597fb86295a Mon Sep 17 00:00:00 2001
From: Alex Pilon <apilon@hashicorp.com>
Date: Tue, 8 Dec 2020 15:40:09 -0500
Subject: [PATCH] clean up internals with typedefs, surface WasInitialized
 error

---
 internal/langserver/diagnostics/diagnostics.go  | 17 +++++++++++------
 .../langserver/diagnostics/diagnostics_test.go  |  4 ++--
 .../langserver/handlers/command/validate.go     |  5 ++++-
 3 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/internal/langserver/diagnostics/diagnostics.go b/internal/langserver/diagnostics/diagnostics.go
index 4a8354a2..fe14c290 100644
--- a/internal/langserver/diagnostics/diagnostics.go
+++ b/internal/langserver/diagnostics/diagnostics.go
@@ -20,13 +20,17 @@ type diagContext struct {
 	diags  []lsp.Diagnostic
 }
 
+type diagnosticSource string
+
+type fileDiagnostics map[diagnosticSource][]lsp.Diagnostic
+
 // Notifier is a type responsible for queueing hcl diagnostics to be converted
 // and sent to the client
 type Notifier struct {
 	logger         *log.Logger
 	sessCtx        context.Context
 	diags          chan diagContext
-	diagsCache     map[string]map[string][]lsp.Diagnostic
+	diagsCache     map[lsp.DocumentURI]fileDiagnostics
 	closeDiagsOnce sync.Once
 }
 
@@ -35,7 +39,7 @@ func NewNotifier(sessCtx context.Context, logger *log.Logger) *Notifier {
 		logger:     logger,
 		sessCtx:    sessCtx,
 		diags:      make(chan diagContext, 50),
-		diagsCache: make(map[string]map[string][]lsp.Diagnostic),
+		diagsCache: make(map[lsp.DocumentURI]fileDiagnostics),
 	}
 	go n.notify()
 	return n
@@ -67,7 +71,7 @@ func (n *Notifier) notify() {
 	for d := range n.diags {
 		if err := jrpc2.PushNotify(d.ctx, "textDocument/publishDiagnostics", lsp.PublishDiagnosticsParams{
 			URI:         d.uri,
-			Diagnostics: n.mergeDiags(string(d.uri), d.source, d.diags),
+			Diagnostics: n.mergeDiags(d.uri, d.source, d.diags),
 		}); err != nil {
 			n.logger.Printf("Error pushing diagnostics: %s", err)
 		}
@@ -77,13 +81,14 @@ func (n *Notifier) notify() {
 // mergeDiags will return all diags from all cached sources for a given uri.
 // the passed diags overwrites the cached entry for the passed source key
 // even if empty
-func (n *Notifier) mergeDiags(uri string, source string, diags []lsp.Diagnostic) []lsp.Diagnostic {
+func (n *Notifier) mergeDiags(uri lsp.DocumentURI, source string, diags []lsp.Diagnostic) []lsp.Diagnostic {
+
 	fileDiags, ok := n.diagsCache[uri]
 	if !ok {
-		fileDiags = make(map[string][]lsp.Diagnostic)
+		fileDiags = make(fileDiagnostics)
 	}
 
-	fileDiags[source] = diags
+	fileDiags[diagnosticSource(source)] = diags
 	n.diagsCache[uri] = fileDiags
 
 	all := []lsp.Diagnostic{}
diff --git a/internal/langserver/diagnostics/diagnostics_test.go b/internal/langserver/diagnostics/diagnostics_test.go
index 0589324e..3812e8ca 100644
--- a/internal/langserver/diagnostics/diagnostics_test.go
+++ b/internal/langserver/diagnostics/diagnostics_test.go
@@ -51,7 +51,7 @@ func TestPublish_DoesNotSendAfterClose(t *testing.T) {
 }
 
 func TestMergeDiags_CachesMultipleSourcesPerURI(t *testing.T) {
-	uri := "test.tf"
+	uri := lsp.DocumentURI("test.tf")
 
 	n := NewNotifier(context.Background(), discardLogger)
 
@@ -77,7 +77,7 @@ func TestMergeDiags_CachesMultipleSourcesPerURI(t *testing.T) {
 }
 
 func TestMergeDiags_OverwritesSource_EvenWithEmptySlice(t *testing.T) {
-	uri := "test.tf"
+	uri := lsp.DocumentURI("test.tf")
 
 	n := NewNotifier(context.Background(), discardLogger)
 
diff --git a/internal/langserver/handlers/command/validate.go b/internal/langserver/handlers/command/validate.go
index 155bf5d3..b215f449 100644
--- a/internal/langserver/handlers/command/validate.go
+++ b/internal/langserver/handlers/command/validate.go
@@ -29,7 +29,10 @@ func TerraformValidateHandler(ctx context.Context, args cmd.CommandArgs) (interf
 		return nil, err
 	}
 
-	wasInit, _ := rm.WasInitialized()
+	wasInit, err := rm.WasInitialized()
+	if err != nil {
+		return nil, fmt.Errorf("error checking if %s was initialized: %s", dirUri, err)
+	}
 	if !wasInit {
 		return nil, fmt.Errorf("%s is not an initialized module, terraform validate cannot be called", dirUri)
 	}