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) { 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 } 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) } } 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) } }