diff --git a/internal/lsp/lsp_test.go b/internal/lsp/lsp_test.go
index b9d0c27..f45290d 100644
--- a/internal/lsp/lsp_test.go
+++ b/internal/lsp/lsp_test.go
@@ -9,6 +9,8 @@
 	"context"
 	"fmt"
 	"go/token"
+	"os"
+	"os/exec"
 	"sort"
 	"strings"
 	"testing"
@@ -226,8 +228,24 @@
 
 func (r *runner) Format(t *testing.T, data tests.Formats) {
 	ctx := context.Background()
-	for filename, gofmted := range data {
-		uri := span.FileURI(filename)
+	for _, spn := range data {
+		uri := spn.URI()
+		filename, err := uri.Filename()
+		if err != nil {
+			t.Fatal(err)
+		}
+		gofmted := string(r.data.Golden("gofmt", filename, func(golden string) error {
+			cmd := exec.Command("gofmt", filename)
+			stdout, err := os.Create(golden)
+			if err != nil {
+				return err
+			}
+			defer stdout.Close()
+			cmd.Stdout = stdout
+			cmd.Run() // ignore error, sometimes we have intentionally ungofmt-able files
+			return nil
+		}))
+
 		edits, err := r.server.Formatting(context.Background(), &protocol.DocumentFormattingParams{
 			TextDocument: protocol.TextDocumentIdentifier{
 				URI: protocol.NewURI(uri),
diff --git a/internal/lsp/tests/tests.go b/internal/lsp/tests/tests.go
index d4a8960..196087d 100644
--- a/internal/lsp/tests/tests.go
+++ b/internal/lsp/tests/tests.go
@@ -11,7 +11,6 @@
 	"go/parser"
 	"go/token"
 	"io/ioutil"
-	"os"
 	"os/exec"
 	"path"
 	"path/filepath"
@@ -50,7 +49,7 @@
 type Diagnostics map[span.URI][]source.Diagnostic
 type CompletionItems map[token.Pos]*source.CompletionItem
 type Completions map[span.Span][]token.Pos
-type Formats map[string]string
+type Formats []span.Span
 type Definitions map[span.Span]Definition
 type Highlights map[string][]span.Span
 type Symbols map[span.URI][]source.Symbol
@@ -100,7 +99,6 @@
 		Diagnostics:     make(Diagnostics),
 		CompletionItems: make(CompletionItems),
 		Completions:     make(Completions),
-		Formats:         make(Formats),
 		Definitions:     make(Definitions),
 		Highlights:      make(Highlights),
 		Symbols:         make(Symbols),
@@ -317,18 +315,8 @@
 	}
 }
 
-func (data *Data) collectFormats(pos token.Position) {
-	data.Formats[pos.Filename] = string(data.Golden("gofmt", pos.Filename, func(golden string) error {
-		cmd := exec.Command("gofmt", pos.Filename)
-		stdout, err := os.Create(golden)
-		if err != nil {
-			return err
-		}
-		defer stdout.Close()
-		cmd.Stdout = stdout
-		cmd.Run() // ignore error, sometimes we have intentionally ungofmt-able files
-		return nil
-	}))
+func (data *Data) collectFormats(spn span.Span) {
+	data.Formats = append(data.Formats, spn)
 }
 
 func (data *Data) collectDefinitions(src, target span.Span) {