petal/internal/suggestions/handlers_test.go

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