package suggestions import ( "bytes" "context" "encoding/json" "net/http" "net/http/httptest" "path/filepath" "testing" "github.com/go-chi/chi/v5" "gitea.parodia.dev/drwily/petal/internal/db" "gitea.parodia.dev/drwily/petal/internal/llm" ) // stubClient returns a canned checkpoint response; it never touches the network. type stubClient struct { response string calls int } func (s *stubClient) Complete(_ context.Context, _ llm.CompletionRequest) (string, error) { s.calls++ return s.response, nil } func (s *stubClient) Stream(_ context.Context, _ llm.CompletionRequest) (<-chan string, error) { ch := make(chan string) close(ch) return ch, nil } // newTestServer wires a real DB, a seeded document, and the suggestion routes // (both the doc-scoped and the per-suggestion mounts) onto one router. func newTestServer(t *testing.T, client llm.LLMClient) (http.Handler, string, *Handler) { t.Helper() database, err := db.Open(filepath.Join(t.TempDir(), "test.db")) if err != nil { t.Fatalf("open db: %v", err) } t.Cleanup(func() { database.Close() }) // Seed a document with some content to check. var docID string err = database.QueryRow( `INSERT INTO documents (user_id, content_text) VALUES (?, ?) RETURNING id`, db.LocalUserID, "I has two apple.", ).Scan(&docID) if err != nil { t.Fatalf("seed doc: %v", err) } h := New(database, client) r := chi.NewRouter() r.Route("/docs", func(dr chi.Router) { h.RegisterDocRoutes(dr) }) r.Mount("/suggestions", h.Routes()) return r, docID, h } func do(t *testing.T, srv http.Handler, method, path, body string) *httptest.ResponseRecorder { t.Helper() var r *http.Request if body != "" { r = httptest.NewRequest(method, path, bytes.NewBufferString(body)) } else { r = httptest.NewRequest(method, path, nil) } rec := httptest.NewRecorder() srv.ServeHTTP(rec, r) return rec } func TestCheckpointFlow(t *testing.T) { client := &stubClient{response: `{"suggestions":[ {"original":"I has","replacement":"I have","explanation":"subject-verb agreement","type":"grammar"}, {"original":"two apple","replacement":"two apples","explanation":"plural noun","type":"grammar"} ]}`} srv, docID, _ := newTestServer(t, client) // Run a checkpoint → two pending suggestions, with positions located. rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "") if rec.Code != http.StatusOK { t.Fatalf("check: code=%d body=%s", rec.Code, rec.Body) } var got []db.Suggestion if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { t.Fatalf("decode: %v", err) } if len(got) != 2 { t.Fatalf("want 2 suggestions, got %d: %+v", len(got), got) } if got[0].Original != "I has" || got[0].FromPos != 0 { t.Fatalf("first suggestion anchoring wrong: %+v", got[0]) } if got[0].Status != db.SuggestionStatusPending { t.Fatalf("new suggestion should be pending: %+v", got[0]) } // Listing returns the same pending set. rec = do(t, srv, http.MethodGet, "/docs/"+docID+"/suggestions", "") var listed []db.Suggestion _ = json.Unmarshal(rec.Body.Bytes(), &listed) if len(listed) != 2 { t.Fatalf("list pending: want 2, got %d", len(listed)) } // Accept the first, dismiss the second. if rec = do(t, srv, http.MethodPost, "/suggestions/"+got[0].ID+"/accept", ""); rec.Code != http.StatusNoContent { t.Fatalf("accept: code=%d body=%s", rec.Code, rec.Body) } if rec = do(t, srv, http.MethodPost, "/suggestions/"+got[1].ID+"/dismiss", ""); rec.Code != http.StatusNoContent { t.Fatalf("dismiss: code=%d body=%s", rec.Code, rec.Body) } // Pending list is now empty (accepted/dismissed are excluded). rec = do(t, srv, http.MethodGet, "/docs/"+docID+"/suggestions", "") _ = json.Unmarshal(rec.Body.Bytes(), &listed) if len(listed) != 0 { t.Fatalf("pending after resolve: want 0, got %d", len(listed)) } // Re-accepting an already-resolved suggestion 404s. if rec = do(t, srv, http.MethodPost, "/suggestions/"+got[0].ID+"/accept", ""); rec.Code != http.StatusNotFound { t.Fatalf("re-accept: want 404, got %d", rec.Code) } } // TestResolvedSuggestionsNotReproposed proves a checkpoint won't re-nag a // sentence the user already accepted or dismissed: the model returns the same // raw batch on the next pass (it has no memory), but the resolved edits are // suppressed. This is the "accept it, then it nags again moments later" bug. func TestResolvedSuggestionsNotReproposed(t *testing.T) { client := &stubClient{response: `{"suggestions":[ {"original":"I has","replacement":"I have","explanation":"subject-verb agreement","type":"grammar"}, {"original":"two apple","replacement":"two apples","explanation":"plural noun","type":"grammar"} ]}`} srv, docID, h := newTestServer(t, client) // Zero the checkpoint floor so the second pass runs instead of being throttled. h.Limit = llm.NewRateLimiter(0) rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "") var got []db.Suggestion _ = json.Unmarshal(rec.Body.Bytes(), &got) if len(got) != 2 { t.Fatalf("first pass: want 2, got %d", len(got)) } // Accept one, dismiss the other. do(t, srv, http.MethodPost, "/suggestions/"+got[0].ID+"/accept", "") do(t, srv, http.MethodPost, "/suggestions/"+got[1].ID+"/dismiss", "") // Second checkpoint returns the identical raw batch — both must be suppressed. rec = do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "") if rec.Code != http.StatusOK { t.Fatalf("second check: code=%d body=%s", rec.Code, rec.Body) } var again []db.Suggestion if err := json.Unmarshal(rec.Body.Bytes(), &again); err != nil { t.Fatalf("decode: %v", err) } if len(again) != 0 { t.Fatalf("resolved suggestions should not be re-proposed, got %d: %+v", len(again), again) } } // TestVoicePassCoexists proves the voice pass and the grammar checkpoint own // disjoint pending families: running one never wipes the other's flags, and // each endpoint returns the unified pending set. func TestVoicePassCoexists(t *testing.T) { // One client serves both passes; the prompt distinguishes them but the stub // just echoes whatever we set before each call. client := &switchClient{} srv, docID, h := newTestServer(t, client) // Zero the voice floor so the test can re-run the pass without waiting out // the real 20s interval; the family-scoping logic is what's under test here. h.VoiceLimit = llm.NewRateLimiter(0) // Grammar checkpoint first → one grammar flag. client.response = `{"suggestions":[{"original":"I has","replacement":"I have","explanation":"agreement","type":"grammar"}]}` rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "") if rec.Code != http.StatusOK { t.Fatalf("check: code=%d body=%s", rec.Code, rec.Body) } // Voice pass next → one voice flag (null replacement). It must NOT remove the // grammar flag; the response is the unified set of both. client.response = `{"suggestions":[{"original":"two apple","replacement":null,"explanation":"sounds more formal","type":"voice"}]}` rec = do(t, srv, http.MethodPost, "/docs/"+docID+"/voice", "") if rec.Code != http.StatusOK { t.Fatalf("voice: code=%d body=%s", rec.Code, rec.Body) } var got []db.Suggestion if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { t.Fatalf("decode: %v", err) } if len(got) != 2 { t.Fatalf("voice response should carry both families, got %d: %+v", len(got), got) } var grammar, voice *db.Suggestion for i := range got { switch got[i].Type { case db.SuggestionTypeGrammar: grammar = &got[i] case db.SuggestionTypeVoice: voice = &got[i] } } if grammar == nil || voice == nil { t.Fatalf("want one grammar + one voice flag, got %+v", got) } if voice.Replacement != "" { t.Fatalf("voice flag should have empty replacement, got %q", voice.Replacement) } // A fresh voice pass that finds nothing replaces only the voice flag; the // grammar one survives. client.response = `{"suggestions":[]}` rec = do(t, srv, http.MethodPost, "/docs/"+docID+"/voice", "") if rec.Code != http.StatusOK { t.Fatalf("second voice: code=%d body=%s", rec.Code, rec.Body) } _ = json.Unmarshal(rec.Body.Bytes(), &got) if len(got) != 1 || got[0].Type != db.SuggestionTypeGrammar { t.Fatalf("after clearing voice, only the grammar flag should remain, got %+v", got) } } // TestCollocationPassCoexists proves the collocation coach is a third // independent family: it never wipes grammar or voice pending flags, a grammar // checkpoint never wipes its flags, and each endpoint returns the unified set. func TestCollocationPassCoexists(t *testing.T) { client := &switchClient{} srv, docID, h := newTestServer(t, client) // Zero the floors so the test can re-run passes without waiting them out. h.Limit = llm.NewRateLimiter(0) h.VoiceLimit = llm.NewRateLimiter(0) h.CollocationLimit = llm.NewRateLimiter(0) // Grammar + voice first, so all three families are exercised. client.response = `{"suggestions":[{"original":"I has","replacement":"I have","explanation":"agreement","type":"grammar"}]}` if rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/check", ""); rec.Code != http.StatusOK { t.Fatalf("check: code=%d body=%s", rec.Code, rec.Body) } client.response = `{"suggestions":[{"original":"two apple","replacement":null,"explanation":"sounds formal","type":"voice"}]}` if rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/voice", ""); rec.Code != http.StatusOK { t.Fatalf("voice: code=%d body=%s", rec.Code, rec.Body) } // Collocation pass → a third flag (with a replacement). It must keep the // grammar and voice flags; the response is the unified set of all three. client.response = `{"suggestions":[{"original":"two apple","replacement":"two apples","explanation":"Natives usually say…","type":"collocation"}]}` rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/collocation", "") if rec.Code != http.StatusOK { t.Fatalf("collocation: code=%d body=%s", rec.Code, rec.Body) } var got []db.Suggestion if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { t.Fatalf("decode: %v", err) } byType := map[string]bool{} for _, s := range got { byType[s.Type] = true } if !byType[db.SuggestionTypeGrammar] || !byType[db.SuggestionTypeVoice] || !byType[db.SuggestionTypeCollocation] { t.Fatalf("collocation response should carry all three families, got %+v", got) } // A grammar checkpoint must NOT wipe the voice or collocation flags. client.response = `{"suggestions":[]}` rec = do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "") if rec.Code != http.StatusOK { t.Fatalf("second check: code=%d body=%s", rec.Code, rec.Body) } _ = json.Unmarshal(rec.Body.Bytes(), &got) byType = map[string]bool{} for _, s := range got { byType[s.Type] = true } if byType[db.SuggestionTypeGrammar] { t.Fatalf("grammar flag should have cleared, got %+v", got) } if !byType[db.SuggestionTypeVoice] || !byType[db.SuggestionTypeCollocation] { t.Fatalf("grammar checkpoint wiped a sibling family, got %+v", got) } } // switchClient returns a response that can be swapped between calls. type switchClient struct { response string calls int } func (s *switchClient) Complete(_ context.Context, _ llm.CompletionRequest) (string, error) { s.calls++ return s.response, nil } func (s *switchClient) Stream(_ context.Context, _ llm.CompletionRequest) (<-chan string, error) { ch := make(chan string) close(ch) return ch, nil } func TestCheckpointRateLimit(t *testing.T) { client := &stubClient{response: `{"suggestions":[{"original":"I has","replacement":"I have","explanation":"x","type":"grammar"}]}`} srv, docID, _ := newTestServer(t, client) // First check runs the model. do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "") // Immediate second check is throttled — returns the existing set without // hitting the model again. do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "") if client.calls != 1 { t.Fatalf("rate limit should suppress the second model call, got %d calls", client.calls) } } // postMechanics submits a deterministic-findings batch (as the client's prose // detector would) and returns the unified pending set. func postMechanics(t *testing.T, srv http.Handler, docID, findingsJSON string) []db.Suggestion { t.Helper() rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/mechanics", `{"findings":`+findingsJSON+`}`) if rec.Code != http.StatusOK { t.Fatalf("mechanics: code=%d body=%s", rec.Code, rec.Body) } var got []db.Suggestion if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { t.Fatalf("decode: %v", err) } return got } // TestMechanicsPersistsFindings proves the endpoint persists client-detected // findings as a 'mechanics' family with the exact offsets the client supplied, // and that a later grammar checkpoint leaves them untouched (own family). func TestMechanicsPersistsFindings(t *testing.T) { client := &stubClient{response: `{"suggestions":[]}`} srv, docID, _ := newTestServer(t, client) got := postMechanics(t, srv, docID, `[ {"from":0,"to":1,"original":"i","replacement":"I","explanation":"capitalize I"}, {"from":6,"to":13,"original":"the the","replacement":"the","explanation":"doubled word"} ]`) var mech []db.Suggestion for _, s := range got { if s.Type == db.SuggestionTypeMechanics { mech = append(mech, s) } } if len(mech) != 2 { t.Fatalf("want 2 mechanics findings, got %d: %+v", len(mech), got) } // The client's exact offsets are preserved verbatim. for _, s := range mech { if s.Original == "the the" && (s.FromPos != 6 || s.ToPos != 13) { t.Errorf("offsets not preserved: %+v", s) } } // A grammar checkpoint must not wipe the mechanics family. rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "") if rec.Code != http.StatusOK { t.Fatalf("check: code=%d body=%s", rec.Code, rec.Body) } var after []db.Suggestion _ = json.Unmarshal(rec.Body.Bytes(), &after) count := 0 for _, s := range after { if s.Type == db.SuggestionTypeMechanics { count++ } } if count != 2 { t.Fatalf("grammar checkpoint disturbed the mechanics family: got %d, want 2", count) } } // TestMechanicsWinsSpanCollision proves that when a mechanics finding and an LLM // grammar suggestion claim overlapping characters, the LLM card is dropped from // the response and the exact mechanics fix owns the span. func TestMechanicsWinsSpanCollision(t *testing.T) { // LLM grammar suggestion covers "the the cat", overlapping the mechanics // "the the" finding at [0,7]. client := &stubClient{response: `{"suggestions":[ {"original":"the the cat","replacement":"the cat","explanation":"wordy","type":"grammar"} ]}`} srv, docID, h := newTestServer(t, client) if _, err := h.DB.Exec(`UPDATE documents SET content_text = ? WHERE id = ?`, "the the cat sat", docID); err != nil { t.Fatalf("set content: %v", err) } // Persist the mechanics finding, then run the grammar pass. postMechanics(t, srv, docID, `[{"from":0,"to":7,"original":"the the","replacement":"the","explanation":"doubled word"}]`) rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "") if rec.Code != http.StatusOK { t.Fatalf("check: code=%d body=%s", rec.Code, rec.Body) } var got []db.Suggestion if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { t.Fatalf("decode: %v", err) } for _, s := range got { if s.Type == db.SuggestionTypeGrammar { t.Fatalf("overlapping grammar card should have been dropped in favor of mechanics, got %+v", got) } } if len(got) != 1 || got[0].Type != db.SuggestionTypeMechanics { t.Fatalf("want sole mechanics finding, got %+v", got) } } // TestMechanicsActionedSuppression proves a dismissed mechanics fix is not // re-proposed on the next submission of the same finding. func TestMechanicsActionedSuppression(t *testing.T) { client := &stubClient{response: `{"suggestions":[]}`} srv, docID, _ := newTestServer(t, client) findings := `[{"from":6,"to":13,"original":"the the","replacement":"the","explanation":"doubled word"}]` got := postMechanics(t, srv, docID, findings) if len(got) != 1 { t.Fatalf("want 1 finding, got %+v", got) } // Dismiss it, then resubmit the same finding — it must stay suppressed. do(t, srv, http.MethodPost, "/suggestions/"+got[0].ID+"/dismiss", "") again := postMechanics(t, srv, docID, findings) if len(again) != 0 { t.Fatalf("dismissed mechanics fix should not reappear, got %+v", again) } }