review-bot/budget/budget_test.go

package budget

import (
	"strings"
	"testing"
)

func TestEstimateTokens(t *testing.T) {
	tests := []struct {
		input string
		want  int
	}{
		{"", 0},
		{"abcd", 1},
		{"12345678", 2},
		{strings.Repeat("x", 400), 100},
	}
	for _, tt := range tests {
		got := EstimateTokens(tt.input)
		if got != tt.want {
			t.Errorf("EstimateTokens(%d chars) = %d, want %d", len(tt.input), got, tt.want)
		}
	}
}

func TestLimitForModel(t *testing.T) {
	tests := []struct {
		model string
		want  int
	}{
		{"gpt-4.1", 128_000},
		{"gpt-5", 200_000},
		{"gpt-5-mini", 200_000},
		{"unknown-model", defaultLimit},
		{"gpt-4.1-2026-01-01", 128_000}, // prefix match
	}
	for _, tt := range tests {
		got := LimitForModel(tt.model)
		if got != tt.want {
			t.Errorf("LimitForModel(%q) = %d, want %d", tt.model, got, tt.want)
		}
	}
}

func TestFit_AllFits(t *testing.T) {
	s := Sections{
		SystemBase:  "system instructions",
		Patterns:    "some patterns",
		Conventions: "some conventions",
		FileContext: "file content",
		Diff:        "diff content",
		UserMeta:    "PR: title\n",
	}
	result := Fit("gpt-5", s)

	if len(result.Trimmed) != 0 {
		t.Errorf("expected no trimming, got %v", result.Trimmed)
	}
	if !strings.Contains(result.SystemPrompt, "some patterns") {
		t.Error("expected patterns in system prompt")
	}
	if !strings.Contains(result.SystemPrompt, "some conventions") {
		t.Error("expected conventions in system prompt")
	}
	if !strings.Contains(result.UserPrompt, "file content") {
		t.Error("expected file context in user prompt")
	}
}

func TestFit_TrimsPatterns(t *testing.T) {
	// Create content that exceeds 128K token budget for gpt-4.1
	// Budget ≈ 128K - 4K reserve = 124K tokens = ~496K chars
	// Fill patterns with enough to push over
	bigPatterns := strings.Repeat("x", 500_000) // ~125K tokens
	s := Sections{
		SystemBase:  "base",
		Patterns:    bigPatterns,
		Conventions: "conventions",
		FileContext: "files",
		Diff:        "diff",
		UserMeta:    "meta",
	}
	result := Fit("gpt-4.1", s)

	if len(result.Trimmed) == 0 {
		t.Fatal("expected trimming")
	}
	if !strings.Contains(result.Trimmed[0], "patterns") {
		t.Errorf("expected patterns to be trimmed first, got %v", result.Trimmed)
	}
	if strings.Contains(result.SystemPrompt, bigPatterns[:100]) {
		t.Error("expected patterns to be removed from output")
	}
	// Conventions should survive
	if !strings.Contains(result.SystemPrompt, "conventions") {
		t.Error("expected conventions to survive after patterns trimmed")
	}
}

func TestFit_TrimsConventions(t *testing.T) {
	// Patterns + conventions + diff all exceed budget even after patterns removed
	big := strings.Repeat("y", 520_000) // ~130K tokens each (exceeds 124K budget even alone)
	s := Sections{
		SystemBase:  "base",
		Patterns:    big,
		Conventions: big,
		FileContext: "files",
		Diff:        "diff",
		UserMeta:    "meta",
	}
	result := Fit("gpt-4.1", s)

	if len(result.Trimmed) < 2 {
		t.Fatalf("expected at least 2 trimmed, got %v", result.Trimmed)
	}
	if !strings.Contains(result.Trimmed[0], "patterns") {
		t.Errorf("expected patterns trimmed first, got %s", result.Trimmed[0])
	}
	if !strings.Contains(result.Trimmed[1], "conventions") {
		t.Errorf("expected conventions trimmed second, got %s", result.Trimmed[1])
	}
}

func TestFit_TruncatesDiff(t *testing.T) {
	// Only diff is huge, no patterns/conventions
	hugeDiff := strings.Repeat("z", 600_000) // ~150K tokens > 128K limit
	s := Sections{
		SystemBase: "base",
		Diff:       hugeDiff,
		UserMeta:   "meta",
	}
	result := Fit("gpt-4.1", s)

	if len(result.Trimmed) == 0 {
		t.Fatal("expected diff truncation")
	}
	if !strings.Contains(result.Trimmed[len(result.Trimmed)-1], "diff truncated") {
		t.Errorf("expected diff truncation note, got %v", result.Trimmed)
	}
	if !strings.Contains(result.UserPrompt, "[diff truncated due to context limit]") {
		t.Error("expected truncation marker in user prompt")
	}
}

func TestFit_PreservesNoteInOutput(t *testing.T) {
	big := strings.Repeat("w", 500_000)
	s := Sections{
		SystemBase: "base",
		Patterns:   big,
		Diff:       "small diff",
		UserMeta:   "meta",
	}
	result := Fit("gpt-4.1", s)

	if !strings.Contains(result.UserPrompt, "⚠️ Note: Context was trimmed") {
		t.Error("expected trimming note in user prompt")
	}
}