petal/internal/llm/checkpoint_test.go

package llm

import (
	"testing"
	"time"
)

func TestParseCheckpoint(t *testing.T) {
	tests := []struct {
		name    string
		raw     string
		want    int
		wantErr bool
	}{
		{
			name: "clean json",
			raw:  `{"suggestions":[{"original":"I has","replacement":"I have","explanation":"subject-verb agreement","type":"grammar"}]}`,
			want: 1,
		},
		{
			name: "wrapped in markdown fence",
			raw:  "Here you go:\n```json\n{\"suggestions\":[{\"original\":\"a apple\",\"replacement\":\"an apple\",\"explanation\":\"use an before a vowel\",\"type\":\"grammar\"}]}\n```",
			want: 1,
		},
		{
			name: "empty suggestions",
			raw:  `{"suggestions":[]}`,
			want: 0,
		},
		{
			name: "drops items with empty original",
			raw:  `{"suggestions":[{"original":"","replacement":"x","explanation":"e","type":"grammar"},{"original":"teh","replacement":"the","explanation":"typo","type":"grammar"}]}`,
			want: 1,
		},
		{
			name:    "no json at all",
			raw:     "I could not find any issues!",
			wantErr: true,
		},
		{
			name: "braces inside string values",
			raw:  `{"suggestions":[{"original":"use {x}","replacement":"use x","explanation":"drop the braces","type":"clarity"}]}`,
			want: 1,
		},
		{
			// The bug that 502'd every long doc: the model hit num_predict and the
			// JSON never closed. Two suggestions completed; the third is cut off.
			// Salvage keeps the two that completed instead of failing outright.
			name: "truncated mid-array salvages completed objects",
			raw:  `{"suggestions":[{"original":"I has","replacement":"I have","explanation":"subject-verb agreement","type":"grammar"},{"original":"a apple","replacement":"an apple","explanation":"use an before a vowel","type":"grammar"},{"original":"she relised","replacement":"she real`,
			want: 2,
		},
		{
			// A truncated array where even the first object is incomplete has
			// nothing to salvage → still an error (the caller releases + retries).
			name:    "truncated before any object closes",
			raw:     `{"suggestions":[{"original":"It was a thursdy morning when Clara`,
			wantErr: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := ParseCheckpoint(tt.raw)
			if tt.wantErr {
				if err == nil {
					t.Fatalf("expected error, got none")
				}
				return
			}
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}
			if len(got) != tt.want {
				t.Fatalf("got %d suggestions, want %d", len(got), tt.want)
			}
		})
	}
}

func TestRateLimiter(t *testing.T) {
	rl := NewRateLimiter(30 * time.Second)

	ok, _, at := rl.Allow("doc1")
	if !ok {
		t.Fatal("first call should be allowed")
	}
	if ok, retry, _ := rl.Allow("doc1"); ok || retry <= 0 {
		t.Fatalf("immediate second call should be throttled, got ok=%v retry=%v", ok, retry)
	}
	// A different document is independent.
	if ok, _, _ := rl.Allow("doc2"); !ok {
		t.Fatal("different doc should be allowed")
	}

	// Releasing doc1's slot (as a failed pass does) lets the next check re-run
	// immediately instead of waiting out the interval.
	rl.Release("doc1", at)
	if ok, _, _ := rl.Allow("doc1"); !ok {
		t.Fatal("after Release the next call should be allowed again")
	}
	// Release only rolls back its own slot: a stale timestamp must not evict the
	// slot a newer Allow holds (so a late failure can't cancel a fresh in-flight
	// check). doc2 still holds its slot from above; a zero-time Release is a no-op.
	rl.Release("doc2", time.Time{})
	if ok, _, _ := rl.Allow("doc2"); ok {
		t.Fatal("stale Release must not free doc2's current slot")
	}
}

func TestTruncateDoc(t *testing.T) {
	short := "hello"
	if got := TruncateDoc(short); got != short {
		t.Fatalf("short doc should be unchanged")
	}
	long := make([]byte, maxDocChars+500)
	for i := range long {
		long[i] = 'a'
	}
	got := TruncateDoc(string(long))
	if len(got) != maxDocChars {
		t.Fatalf("truncated length = %d, want %d", len(got), maxDocChars)
	}
}