PgTidy/pkg/lsp/server_test.go

package lsp

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"testing"
	"time"
)

// rpc sends a JSON-RPC request frame and returns the raw body bytes.
func frame(id int, method string, params interface{}) []byte {
	type req struct {
		JSONRPC string      `json:"jsonrpc"`
		ID      int         `json:"id"`
		Method  string      `json:"method"`
		Params  interface{} `json:"params"`
	}
	body, _ := json.Marshal(req{JSONRPC: "2.0", ID: id, Method: method, Params: params})
	return []byte(fmt.Sprintf("Content-Length: %d\r\n\r\n%s", len(body), body))
}

func notifFrame(method string, params interface{}) []byte {
	type notif struct {
		JSONRPC string      `json:"jsonrpc"`
		Method  string      `json:"method"`
		Params  interface{} `json:"params"`
	}
	body, _ := json.Marshal(notif{JSONRPC: "2.0", Method: method, Params: params})
	return []byte(fmt.Sprintf("Content-Length: %d\r\n\r\n%s", len(body), body))
}

// readResp reads one JSON-RPC response from a *bytes.Buffer (blocking until available).
func readResp(t *testing.T, buf *bytes.Buffer) map[string]interface{} {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		data := buf.Bytes()
		// Find Content-Length header
		idx := bytes.Index(data, []byte("Content-Length: "))
		if idx < 0 {
			time.Sleep(5 * time.Millisecond)
			continue
		}
		eol := bytes.Index(data[idx:], []byte("\r\n"))
		if eol < 0 {
			time.Sleep(5 * time.Millisecond)
			continue
		}
		lenStr := string(data[idx+16 : idx+eol])
		var n int
		fmt.Sscanf(lenStr, "%d", &n)
		sep := bytes.Index(data, []byte("\r\n\r\n"))
		if sep < 0 || len(data) < sep+4+n {
			time.Sleep(5 * time.Millisecond)
			continue
		}
		body := data[sep+4 : sep+4+n]
		buf.Next(sep + 4 + n)
		var result map[string]interface{}
		json.Unmarshal(body, &result)
		return result
	}
	t.Fatal("timeout waiting for response")
	return nil
}

func runServer(t *testing.T, input []byte) *bytes.Buffer {
	t.Helper()
	out := &bytes.Buffer{}
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	done := make(chan error, 1)
	go func() {
		done <- Serve(ctx, bytes.NewReader(input), out, t.TempDir())
		cancel()
	}()
	// Give the server time to process
	select {
	case <-done:
	case <-time.After(3 * time.Second):
	}
	return out
}

func TestInitialize(t *testing.T) {
	input := append(frame(1, "initialize", map[string]interface{}{}),
		notifFrame("initialized", map[string]interface{}{})...)
	input = append(input, frame(2, "shutdown", nil)...)
	input = append(input, notifFrame("exit", nil)...)

	out := runServer(t, input)

	resp := readResp(t, out)
	if resp["id"].(float64) != 1 {
		t.Fatalf("expected id=1, got %v", resp["id"])
	}
	caps := resp["result"].(map[string]interface{})["capabilities"].(map[string]interface{})
	if caps["documentFormattingProvider"] != true {
		t.Error("expected documentFormattingProvider=true")
	}
}

func TestFormatting(t *testing.T) {
	unformatted := "create function foo() returns void language plpgsql as $$ begin end $$;"
	uri := "file:///test.sql"

	var input []byte
	input = append(input, frame(1, "initialize", map[string]interface{}{})...)
	input = append(input, notifFrame("initialized", nil)...)
	input = append(input, notifFrame("textDocument/didOpen", map[string]interface{}{
		"textDocument": map[string]interface{}{
			"uri":        uri,
			"languageId": "sql",
			"version":    1,
			"text":       unformatted,
		},
	})...)
	input = append(input, frame(2, "textDocument/formatting", map[string]interface{}{
		"textDocument": map[string]interface{}{"uri": uri},
		"options":      map[string]interface{}{"tabSize": 2, "insertSpaces": true},
	})...)
	input = append(input, frame(3, "shutdown", nil)...)
	input = append(input, notifFrame("exit", nil)...)

	out := runServer(t, input)

	// Skip initialize response and publishDiagnostics notification, find formatting response
	var formattingResp map[string]interface{}
	for i := 0; i < 10; i++ {
		resp := readResp(t, out)
		if resp == nil {
			break
		}
		id, hasID := resp["id"]
		if hasID && id.(float64) == 2 {
			formattingResp = resp
			break
		}
	}
	if formattingResp == nil {
		t.Fatal("did not receive formatting response")
	}
	result, ok := formattingResp["result"].([]interface{})
	if !ok {
		t.Fatalf("expected array result, got %T: %v", formattingResp["result"], formattingResp["result"])
	}
	if len(result) == 0 {
		t.Fatal("expected at least one text edit")
	}
	edit := result[0].(map[string]interface{})
	newText := edit["newText"].(string)
	if !strings.Contains(newText, "CREATE FUNCTION") {
		t.Errorf("formatted output missing CREATE FUNCTION keyword; got:\n%s", newText)
	}
}

func TestDidClose_ClearsDiagnostics(t *testing.T) {
	uri := "file:///test.sql"
	// SQL with a lint violation (SELECT * triggers COR001)
	sql := "SELECT * FROM users;"

	var input []byte
	input = append(input, frame(1, "initialize", map[string]interface{}{})...)
	input = append(input, notifFrame("initialized", nil)...)
	input = append(input, notifFrame("textDocument/didOpen", map[string]interface{}{
		"textDocument": map[string]interface{}{
			"uri": uri, "languageId": "sql", "version": 1, "text": sql,
		},
	})...)
	input = append(input, notifFrame("textDocument/didClose", map[string]interface{}{
		"textDocument": map[string]interface{}{"uri": uri},
	})...)
	input = append(input, frame(2, "shutdown", nil)...)
	input = append(input, notifFrame("exit", nil)...)

	out := runServer(t, input)

	// Collect all publishDiagnostics notifications; the last one for this URI must be empty.
	var lastDiags []interface{}
	for i := 0; i < 20; i++ {
		if out.Len() == 0 {
			break
		}
		resp := readResp(t, out)
		if resp == nil {
			break
		}
		if resp["method"] == "textDocument/publishDiagnostics" {
			p := resp["params"].(map[string]interface{})
			if p["uri"] == uri {
				lastDiags = p["diagnostics"].([]interface{})
			}
		}
	}
	if len(lastDiags) != 0 {
		t.Errorf("expected empty diagnostics after didClose, got %d", len(lastDiags))
	}
}