Use test context

Author: jakebailey

cmd/tsgo/lsp.go

@@ -1,10 +1,13 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"os"
+	"os/signal"
 	"runtime"
+	"syscall"
 
 	"github.com/microsoft/typescript-go/internal/bundled"
 	"github.com/microsoft/typescript-go/internal/core"
@@ -37,6 +40,9 @@ func runLSP(args []string) int {
 		defer profileSession.Stop()
 	}
 
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer stop()
+
 	fs := bundled.WrapFS(osvfs.FS())
 	defaultLibraryPath := bundled.LibPath()
 	typingsLocation := getGlobalTypingsCacheLocation()
@@ -51,7 +57,7 @@ func runLSP(args []string) int {
 		TypingsLocation:    typingsLocation,
 	})
 
-	if err := s.Run(); err != nil {
+	if err := s.Run(ctx); err != nil {
 		return 1
 	}
 	return 0

internal/fourslash/fourslash.go

@@ -175,7 +175,7 @@ func NewFourslash(t *testing.T, capabilities *lsproto.ClientCapabilities, conten
 		defer func() {
 			outputWriter.Close()
 		}()
-		err := server.Run()
+		err := server.Run(t.Context())
 		if err != nil {
 			t.Error("server error:", err)
 		}

internal/lsp/server.go

@@ -5,13 +5,10 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"os"
-	"os/signal"
 	"runtime/debug"
 	"slices"
 	"sync"
 	"sync/atomic"
-	"syscall"
 
 	"github.com/go-json-experiment/json"
 	"github.com/microsoft/typescript-go/internal/collections"
@@ -249,10 +246,7 @@ func (s *Server) RefreshDiagnostics(ctx context.Context) error {
 	return nil
 }
 
-func (s *Server) Run() error {
-	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
-	defer stop()
-
+func (s *Server) Run(ctx context.Context) error {
 	g, ctx := errgroup.WithContext(ctx)
 	g.Go(func() error { return s.dispatchLoop(ctx) })
 	g.Go(func() error { return s.writeLoop(ctx) })