tests: add shared test fixtures, expand provider breadth (multi-choice, error bodies), add LSP rewrite/diagnostics realism and table-driven tests

10 files changed, 282 insertions, 7 deletions

diff --git a/internal/llm/copilot_http_test.go b/internal/llm/copilot_http_test.go
index c029a65..4c2b7fe 100644
--- a/internal/llm/copilot_http_test.go
+++ b/internal/llm/copilot_http_test.go
@@ -6,9 +6,9 @@ import (
     "io"
     "net/http"
     "net/http/httptest"
+    "strings"
     "testing"
     "time"
-    "strings"
     "encoding/base64"
 )
 
@@ -72,6 +72,42 @@ func TestCopilot_CodeCompletion_Success(t *testing.T) {
     }
 }
 
+func TestCopilot_Chat_MultiChoice_And_ErrorBody(t *testing.T) {
+    // Chat multi-choice: return two choices; client returns first content
+    srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+        _ = json.NewEncoder(w).Encode(map[string]any{
+            "choices": []map[string]any{
+                {"index": 0, "finish_reason": "stop", "message": map[string]string{"role": "assistant", "content": "FIRST"}},
+                {"index": 1, "finish_reason": "length", "message": map[string]string{"role": "assistant", "content": "SECOND"}},
+            },
+        })
+    }))
+    defer srv.Close()
+    c := newCopilot(srv.URL, "gpt-4o-mini", "KEY", f64p(0.1)).(copilotClient)
+    // Token success
+    tr := rtFunc2(func(r *http.Request) (*http.Response, error) {
+        if r.URL.Host == "api.github.com" && r.URL.Path == "/copilot_internal/v2/token" {
+            rw := httptest.NewRecorder(); _ = json.NewEncoder(rw).Encode(map[string]string{"token":"tok"}); res := rw.Result(); res.StatusCode = 200; return res, nil
+        }
+        return http.DefaultTransport.RoundTrip(r)
+    })
+    c.httpClient = &http.Client{Transport: tr, Timeout: 5 * time.Second}
+    out, err := c.Chat(context.Background(), []Message{{Role: "user", Content: "hi"}})
+    if err != nil || out != "FIRST" { t.Fatalf("copilot multi-choice: %v %q", err, out) }
+
+    // Non-2xx with error body
+    srv2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+        w.WriteHeader(403)
+        _ = json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{"message":"denied","type":"forbidden"}})
+    }))
+    defer srv2.Close()
+    c2 := newCopilot(srv2.URL, "gpt-4o-mini", "KEY", f64p(0.1)).(copilotClient)
+    c2.httpClient = &http.Client{Transport: tr, Timeout: 5 * time.Second}
+    if _, err := c2.Chat(context.Background(), []Message{{Role:"user", Content:"hi"}}); err == nil {
+        t.Fatalf("expected error for copilot non-2xx with error body")
+    }
+}
+
 func TestParseJWTExp_AndParseInt64(t *testing.T) {
     // Valid base64 payload
     payload := `{"exp": 1700000000}`
diff --git a/internal/llm/openai_http_test.go b/internal/llm/openai_http_test.go
index 7ae34be..78830ba 100644
--- a/internal/llm/openai_http_test.go
+++ b/internal/llm/openai_http_test.go
@@ -47,3 +47,47 @@ func TestHandleOpenAINon2xx_NoErrorBody(t *testing.T) {
     resp := &http.Response{StatusCode: 500, Body: io.NopCloser(strings.NewReader("{}"))}
     if err := handleOpenAINon2xx(resp, time.Now()); err == nil { t.Fatalf("expected http error") }
 }
+
+func TestOpenAI_ChatStream_SSE_ErrorChunk(t *testing.T) {
+    srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+        w.Header().Set("Content-Type", "text/event-stream")
+        io.WriteString(w, "data: {\"error\":{\"message\":\"oops\"}}\n\n")
+        io.WriteString(w, "data: [DONE]\n")
+    }))
+    defer srv.Close()
+    c := newOpenAI(srv.URL, "g", "KEY", f64p(0.2)).(openAIClient)
+    c.httpClient = srv.Client()
+    var got string
+    if err := c.ChatStream(context.Background(), []Message{{Role:"user", Content:"hi"}}, func(s string){ got += s }); err == nil {
+        t.Fatalf("expected error due to error chunk")
+    }
+}
+
+func TestOpenAI_Chat_MultiChoiceAndErrorBody(t *testing.T) {
+    // Multi-choice success: return two choices with different finish reasons
+    srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+        _ = json.NewEncoder(w).Encode(map[string]any{
+            "choices": []map[string]any{
+                {"index": 0, "finish_reason": "stop", "message": map[string]string{"role": "assistant", "content": "FIRST"}},
+                {"index": 1, "finish_reason": "length", "message": map[string]string{"role": "assistant", "content": "SECOND"}},
+            },
+        })
+    }))
+    defer srv.Close()
+    c := newOpenAI(srv.URL, "g", "KEY", f64p(0.2)).(openAIClient)
+    c.httpClient = srv.Client()
+    out, err := c.Chat(context.Background(), []Message{{Role: "user", Content: "hi"}})
+    if err != nil || out != "FIRST" { t.Fatalf("openai multi-choice: %v %q", err, out) }
+
+    // Error body case: non-2xx with error message
+    srv2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+        w.WriteHeader(400)
+        _ = json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{"message": "bad", "type": "invalid"}})
+    }))
+    defer srv2.Close()
+    c2 := newOpenAI(srv2.URL, "g", "KEY", f64p(0.2)).(openAIClient)
+    c2.httpClient = srv2.Client()
+    if _, err := c2.Chat(context.Background(), []Message{{Role: "user", Content: "hi"}}); err == nil {
+        t.Fatalf("expected error from non-2xx with error body")
+    }
+}
diff --git a/internal/lsp/codeaction_more_test.go b/internal/lsp/codeaction_more_test.go
index 387afb5..412d988 100644
--- a/internal/lsp/codeaction_more_test.go
+++ b/internal/lsp/codeaction_more_test.go
@@ -5,11 +5,12 @@ import (
     "path/filepath"
     "strings"
     "testing"
+    tut "codeberg.org/snonux/hexai/internal/testutil"
 )
 
 func TestBuildDocumentCodeAction_AndResolve(t *testing.T) {
     s := newTestServer()
-    s.llmClient = fakeLLM{resp: "// doc\nfunc add(a,b int) int { return a+b }"}
+    s.llmClient = fakeLLM{resp: tut.MultilineDocBlock()+"\n"+"func add(a,b int) int { return a+b }"}
     uri := "file:///doc.go"
     s.setDocument(uri, "package x\nfunc add(a,b int) int {return a+b}")
     p := CodeActionParams{TextDocument: TextDocumentIdentifier{URI: uri}, Range: Range{Start: Position{Line:1, Character:0}, End: Position{Line:1, Character:10}}}
diff --git a/internal/lsp/completion_prefix_strip_test.go b/internal/lsp/completion_prefix_strip_test.go
index 64cca49..99a08d6 100644
--- a/internal/lsp/completion_prefix_strip_test.go
+++ b/internal/lsp/completion_prefix_strip_test.go
@@ -1,8 +1,9 @@
 package lsp
 
 import (
-	"encoding/json"
-	"testing"
+    "encoding/json"
+    "testing"
+    tut "codeberg.org/snonux/hexai/internal/testutil"
 )
 
 func TestStripDuplicateGeneralPrefix_ExactOverlap(t *testing.T) {
@@ -40,7 +41,7 @@ func TestStripDuplicateAssignmentPrefix_AssignAndWalrus(t *testing.T) {
 
 func TestTryLLMCompletion_ManualInvokeAfterWhitespace_Allows(t *testing.T) {
 	s := &Server{maxTokens: 32, triggerChars: []string{".", ":", "/", "_"}, compCache: make(map[string]string)}
-	s.llmClient = fakeLLM{resp: "() *CustData"}
+    s.llmClient = fakeLLM{resp: tut.MultilineFunctionSuggestion()}
 	line := "func fib(i int) " // cursor after space
 	p := CompletionParams{Position: Position{Line: 0, Character: len(line)}, TextDocument: TextDocumentIdentifier{URI: "file://x.go"}}
 	// Simulate manual user invocation (TriggerKind=1)
diff --git a/internal/lsp/handlers_end_to_end_test.go b/internal/lsp/handlers_end_to_end_test.go
index ba4a0bc..73478e9 100644
--- a/internal/lsp/handlers_end_to_end_test.go
+++ b/internal/lsp/handlers_end_to_end_test.go
@@ -9,6 +9,7 @@ import (
     "strings"
     "testing"
     "time"
+    tut "codeberg.org/snonux/hexai/internal/testutil"
 )
 
 // captureResponse decodes a single LSP Response from the server's output buffer.
@@ -190,7 +191,7 @@ func TestHandle_Dispatch_Initialize(t *testing.T) {
 func TestDetectAndHandleChat_InsertsReply(t *testing.T) {
     var out bytes.Buffer
     s := &Server{logger: log.New(io.Discard, "", 0), docs: make(map[string]*document), out: &out}
-    s.llmClient = fakeLLM{resp: "Hello"}
+    s.llmClient = fakeLLM{resp: tut.MultilineChatReply()}
     uri := "file:///chat.go"
     // Place a prompt line with a supported trigger at EOL, then a blank line
     s.setDocument(uri, "What time?>\n\n")
@@ -208,7 +209,9 @@ func TestDetectAndHandleChat_InsertsReply(t *testing.T) {
     if len(we.Changes) == 0 { t.Fatalf("expected changes in edit") }
     edits := we.Changes[uri]
     if len(edits) != 2 { t.Fatalf("expected 2 edits (delete+insert), got %d", len(edits)) }
-    if !strings.Contains(edits[1].NewText, "> Hello") { t.Fatalf("expected reply insertion with '> Hello', got %q", edits[1].NewText) }
+    if !strings.Contains(edits[1].NewText, "> Hello") || !strings.Contains(edits[1].NewText, "multi-line reply") {
+        t.Fatalf("expected multi-line reply insertion, got %q", edits[1].NewText)
+    }
 }
 
 func TestHandleCodeActionResolve_Diagnostics(t *testing.T) {
diff --git a/internal/lsp/init_and_trigger_test.go b/internal/lsp/init_and_trigger_test.go
new file mode 100644
index 0000000..cdc907e
--- /dev/null
+++ b/internal/lsp/init_and_trigger_test.go
@@ -0,0 +1,52 @@
+package lsp
+
+import (
+    "bytes"
+    "encoding/json"
+    "io"
+    "log"
+    "testing"
+)
+
+func TestHandleInitialize_Capabilities(t *testing.T) {
+    var out bytes.Buffer
+    s := &Server{logger: log.New(io.Discard, "", 0), docs: make(map[string]*document), out: &out}
+    s.triggerChars = []string{".", ":"}
+    req := Request{JSONRPC: "2.0", ID: json.RawMessage("7"), Method: "initialize"}
+    out.Reset()
+    s.handleInitialize(req)
+    resp := captureResponse(t, &out)
+    var init InitializeResult
+    b, _ := json.Marshal(resp.Result)
+    if err := json.Unmarshal(b, &init); err != nil { t.Fatalf("decode init: %v", err) }
+    if init.Capabilities.CodeActionProvider == nil { t.Fatalf("expected codeActionProvider") }
+    // CodeActionProvider is any; re-marshal to struct
+    var cap struct{ ResolveProvider bool `json:"resolveProvider"` }
+    cb, _ := json.Marshal(init.Capabilities.CodeActionProvider)
+    _ = json.Unmarshal(cb, &cap)
+    if !cap.ResolveProvider { t.Fatalf("expected resolveProvider=true") }
+    if init.Capabilities.CompletionProvider == nil || len(init.Capabilities.CompletionProvider.TriggerCharacters) == 0 {
+        t.Fatalf("expected trigger characters") }
+}
+
+func TestIsTriggerEvent_Variants(t *testing.T) {
+    s := newTestServer()
+    s.triggerChars = []string{".", ":"}
+    // 1) Manual invoke via context
+    ctx := struct{ TriggerKind int `json:"triggerKind"` }{TriggerKind:1}
+    raw, _ := json.Marshal(ctx)
+    p := CompletionParams{Position: Position{Line:0, Character:1}, Context: json.RawMessage(raw)}
+    if !s.isTriggerEvent(p, "a") { t.Fatalf("manual invoke should trigger") }
+    // 2) TriggerCharacter present and allowed
+    ctx2 := struct{ TriggerKind int `json:"triggerKind"`; TriggerCharacter string `json:"triggerCharacter"` }{TriggerKind:2, TriggerCharacter: "."}
+    raw2, _ := json.Marshal(ctx2)
+    p2 := CompletionParams{Position: Position{Line:0, Character:1}, Context: json.RawMessage(raw2)}
+    if !s.isTriggerEvent(p2, "a.") { t.Fatalf("trigger char should trigger") }
+    // 3) Fallback char left of cursor
+    p3 := CompletionParams{Position: Position{Line:0, Character:3}}
+    if !s.isTriggerEvent(p3, "ab:") { t.Fatalf("fallback char should trigger") }
+    // 4) Bare ';;' disables trigger
+    p4 := CompletionParams{Position: Position{Line:0, Character:2}}
+    if s.isTriggerEvent(p4, ";;") { t.Fatalf("bare ;; should not trigger") }
+}
+
diff --git a/internal/lsp/instruction_table_test.go b/internal/lsp/instruction_table_test.go
new file mode 100644
index 0000000..e92ffde
--- /dev/null
+++ b/internal/lsp/instruction_table_test.go
@@ -0,0 +1,25 @@
+package lsp
+
+import "testing"
+
+func TestFindFirstInstructionInLine_Table(t *testing.T) {
+    cases := []struct{
+        name string
+        line string
+        instr string
+    }{
+        {"strict_semicolon", ";do; trailing", "do"},
+        {"c_block", "x /* add docs */ y", "add docs"},
+        {"html_comment", "<!-- fix --> code", "fix"},
+        {"slash_slash", "code // please refactor", "please refactor"},
+        {"hash", "# summarize", "summarize"},
+        {"double_dash", "-- rewrite quickly", "rewrite quickly"},
+    }
+    for _, c := range cases {
+        instr, _, ok := findFirstInstructionInLine(c.line)
+        if !ok || instr != c.instr {
+            t.Fatalf("%s: got %q ok=%v", c.name, instr, ok)
+        }
+    }
+}
+
diff --git a/internal/lsp/prefix_table_test.go b/internal/lsp/prefix_table_test.go
new file mode 100644
index 0000000..0ca23d2
--- /dev/null
+++ b/internal/lsp/prefix_table_test.go
@@ -0,0 +1,24 @@
+package lsp
+
+import "testing"
+
+func TestPrefixStripping_Table(t *testing.T) {
+    cases := []struct{ name, prefix, sugg, want string }{
+        {"assign_walrus", "name := ", "name := compute()", "compute()"},
+        {"assign_equals", "x = ", "x = y+1", "y+1"},
+        {"general_db", "db.", "db.Query()", "Query()"},
+        {"general_func", "func New ", "func New() *T", "() *T"},
+    }
+    for _, c := range cases {
+        var got string
+        if c.name == "assign_walrus" || c.name == "assign_equals" {
+            got = stripDuplicateAssignmentPrefix(c.prefix, c.sugg)
+        } else {
+            got = stripDuplicateGeneralPrefix(c.prefix, c.sugg)
+        }
+        if got != c.want {
+            t.Fatalf("%s: got %q want %q", c.name, got, c.want)
+        }
+    }
+}
+
diff --git a/internal/lsp/rewrite_diagnostics_realism_test.go b/internal/lsp/rewrite_diagnostics_realism_test.go
new file mode 100644
index 0000000..87ff571
--- /dev/null
+++ b/internal/lsp/rewrite_diagnostics_realism_test.go
@@ -0,0 +1,62 @@
+package lsp
+
+import (
+    "encoding/json"
+    "testing"
+)
+
+func TestResolveRewrite_MultiLine_PreservesRange(t *testing.T) {
+    s := newTestServer()
+    s.llmClient = fakeLLM{resp: "line1\nline2"}
+    uri := "file:///x.go"
+    s.setDocument(uri, "package p\nvar a=1\n")
+    r := Range{Start: Position{Line:1, Character:0}, End: Position{Line:1, Character:5}}
+    payload := struct {
+        Type        string `json:"type"`
+        URI         string `json:"uri"`
+        Range       Range  `json:"range"`
+        Instruction string `json:"instruction"`
+        Selection   string `json:"selection"`
+    }{Type: "rewrite", URI: uri, Range: r, Instruction: "expand", Selection: "var a"}
+    raw, _ := json.Marshal(payload)
+    ca := CodeAction{Title: "Hexai: rewrite selection", Data: raw}
+    resolved, ok := s.resolveCodeAction(ca)
+    if !ok || resolved.Edit == nil { t.Fatalf("expected resolved rewrite edit") }
+    edits := resolved.Edit.Changes[uri]
+    if len(edits) != 1 { t.Fatalf("expected 1 edit") }
+    if edits[0].Range != r { t.Fatalf("range mismatch: got %+v want %+v", edits[0].Range, r) }
+    if edits[0].NewText == "" || !containsNewline(edits[0].NewText) {
+        t.Fatalf("expected multi-line replacement text, got %q", edits[0].NewText)
+    }
+}
+
+func TestResolveDiagnostics_MultiLine_PreservesRange(t *testing.T) {
+    s := newTestServer()
+    s.llmClient = fakeLLM{resp: "fixed\nvalue"}
+    uri := "file:///x.go"
+    s.setDocument(uri, "package p\nvar x = 1\n")
+    r := Range{Start: Position{Line:1, Character:0}, End: Position{Line:1, Character:10}}
+    payload := struct {
+        Type        string       `json:"type"`
+        URI         string       `json:"uri"`
+        Range       Range        `json:"range"`
+        Selection   string       `json:"selection"`
+        Diagnostics []Diagnostic `json:"diagnostics"`
+    }{Type: "diagnostics", URI: uri, Range: r, Selection: "var x = 1", Diagnostics: []Diagnostic{{Range: Range{Start: Position{Line:1}, End: Position{Line:1, Character:5}}, Message: "msg"}}}
+    raw, _ := json.Marshal(payload)
+    ca := CodeAction{Title: "Hexai: resolve diagnostics", Data: raw}
+    resolved, ok := s.resolveCodeAction(ca)
+    if !ok || resolved.Edit == nil { t.Fatalf("expected resolved diagnostics edit") }
+    edits := resolved.Edit.Changes[uri]
+    if len(edits) != 1 { t.Fatalf("expected 1 edit") }
+    if edits[0].Range != r { t.Fatalf("range mismatch: got %+v want %+v", edits[0].Range, r) }
+    if edits[0].NewText == "" || !containsNewline(edits[0].NewText) {
+        t.Fatalf("expected multi-line replacement text, got %q", edits[0].NewText)
+    }
+}
+
+func containsNewline(s string) bool {
+    for i := 0; i < len(s); i++ { if s[i] == '\n' { return true } }
+    return false
+}
+
diff --git a/internal/testutil/fixtures.go b/internal/testutil/fixtures.go
new file mode 100644
index 0000000..41993d3
--- /dev/null
+++ b/internal/testutil/fixtures.go
@@ -0,0 +1,27 @@
+package testutil
+
+// MultilineDocBlock returns a realistic multi-line documentation block.
+func MultilineDocBlock() string {
+    return "// add adds two numbers\n// returns their sum"
+}
+
+// MultilineChatReply returns a multi-line assistant reply for chat tests.
+func MultilineChatReply() string {
+    return "Hello, world!\nThis is a multi-line reply."
+}
+
+// MultilineFunctionSuggestion returns a more realistic multi-line function body suggestion.
+func MultilineFunctionSuggestion() string {
+    return "(ctx context.Context, input string) (*CustData, error) {\n    // TODO: implement\n    return &CustData{}, nil\n}"
+}
+
+// MarkdownCodeFence returns a fenced markdown snippet used in post-processing tests.
+func MarkdownCodeFence() string {
+    return "```go\nname := value\n```"
+}
+
+// MalformedJSON returns a deliberately malformed JSON string.
+func MalformedJSON() string {
+    return "{\"choices\":[{\"delta\":{\"content\":\"oops\"}}]"
+}
+