cliproxy-plugin-commandcode/plugin/usage_test.go

package plugin

import (
	"context"
	"encoding/json"
	"math"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

func TestFlexibleTime(t *testing.T) {
	tests := []struct {
		name     string
		input    string
		wantYear int
	}{
		{
			name:     "unix seconds",
			input:    `{"resetAt": 1741123456}`,
			wantYear: 2025,
		},
		{
			name:     "unix milliseconds",
			input:    `{"resetAt": 1741123456000}`,
			wantYear: 2025,
		},
		{
			name:     "RFC3339 string",
			input:    `{"resetAt": "2025-06-15T12:00:00Z"}`,
			wantYear: 2025,
		},
		{
			name:     "null",
			input:    `{"resetAt": null}`,
			wantYear: 1, // Zero time year
		},
		{
			name:     "empty string",
			input:    `{"resetAt": ""}`,
			wantYear: 1,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var res struct {
				ResetAt FlexibleTime `json:"resetAt"`
			}
			if err := json.Unmarshal([]byte(tt.input), &res); err != nil {
				t.Fatalf("Unmarshal error: %v", err)
			}
			if res.ResetAt.Year() != tt.wantYear {
				t.Errorf("Year = %d, want %d", res.ResetAt.Year(), tt.wantYear)
			}
		})
	}
}

func TestParseAndFormatUsage(t *testing.T) {
	raw := []byte(`{
		"credits": {
			"monthlyCredits": 1000.0,
			"opensourceMonthlyCredits": 500.0,
			"extraBonus": 50.0
		},
		"windowLimits": {
			"fiveHour": {
				"used": 25.0,
				"cap": 100.0,
				"exceeded": false,
				"resetAt": 1741123456
			},
			"weekly": {
				"used": 200.0,
				"cap": 1000.0,
				"exceeded": false,
				"resetAt": "2025-03-10T12:00:00Z"
			}
		}
	}`)

	now := time.Date(2025, 3, 4, 12, 0, 0, 0, time.UTC)

	t.Run("without summary monthly defaults empty", func(t *testing.T) {
		usage, err := ParseAndFormatUsage(raw, nil, now)
		if err != nil {
			t.Fatalf("ParseAndFormatUsage error: %v", err)
		}
		if !usage.OK {
			t.Fatal("expected OK=true")
		}
		if usage.WindowLimits.Monthly.Used != 0 || usage.WindowLimits.Monthly.Cap != 0 {
			t.Errorf("expected empty monthly without summary, got %+v", usage.WindowLimits.Monthly)
		}
	})

	t.Run("monthly window derived from summary + remaining credits", func(t *testing.T) {
		summary := &UpstreamUsageSummaryResponse{TotalMonthlyCredits: 30.0}
		// monthlyCredits in raw = 1000 remaining, so cap = 1030
		usage, err := ParseAndFormatUsage(raw, summary, now)
		if err != nil {
			t.Fatalf("ParseAndFormatUsage error: %v", err)
		}
		m := usage.WindowLimits.Monthly
		if m.Used != 30.0 {
			t.Errorf("monthly used = %v, want 30", m.Used)
		}
		if m.Cap != 1030.0 {
			t.Errorf("monthly cap = %v, want 1030", m.Cap)
		}
		if m.Remaining != 1000.0 {
			t.Errorf("monthly remaining = %v, want 1000", m.Remaining)
		}
		wantPct := math.Round((30.0/1030.0)*10000) / 100
		if m.Percentage != wantPct {
			t.Errorf("monthly percentage = %v, want %v", m.Percentage, wantPct)
		}
	})

	usage, err := ParseAndFormatUsage(raw, nil, now)
	if err != nil {
		t.Fatalf("ParseAndFormatUsage error: %v", err)
	}

	if !usage.OK {
		t.Fatal("expected OK=true")
	}

	// Verify credits
	if usage.Credits.MonthlyCredits != 1000.0 {
		t.Errorf("MonthlyCredits = %v, want 1000", usage.Credits.MonthlyCredits)
	}
	if usage.Credits.OpensourceMonthlyCredits != 500.0 {
		t.Errorf("OpensourceMonthlyCredits = %v, want 500", usage.Credits.OpensourceMonthlyCredits)
	}
	if usage.Credits.TotalCredits != 1500.0 {
		t.Errorf("TotalCredits = %v, want 1500", usage.Credits.TotalCredits)
	}

	// Verify 5-hour window
	fiveHour := usage.WindowLimits.FiveHour
	if fiveHour.Used != 25.0 {
		t.Errorf("FiveHour Used = %v, want 25", fiveHour.Used)
	}
	if fiveHour.Cap != 100.0 {
		t.Errorf("FiveHour Cap = %v, want 100", fiveHour.Cap)
	}
	if fiveHour.Remaining != 75.0 {
		t.Errorf("FiveHour Remaining = %v, want 75", fiveHour.Remaining)
	}
	if fiveHour.Percentage != 25.0 {
		t.Errorf("FiveHour Percentage = %v, want 25", fiveHour.Percentage)
	}
	if fiveHour.Exceeded {
		t.Errorf("FiveHour Exceeded = true, want false")
	}

	// Verify weekly window
	weekly := usage.WindowLimits.Weekly
	if weekly.Used != 200.0 {
		t.Errorf("Weekly Used = %v, want 200", weekly.Used)
	}
	if weekly.Cap != 1000.0 {
		t.Errorf("Weekly Cap = %v, want 1000", weekly.Cap)
	}
	if weekly.Remaining != 800.0 {
		t.Errorf("Weekly Remaining = %v, want 800", weekly.Remaining)
	}
	if weekly.Percentage != 20.0 {
		t.Errorf("Weekly Percentage = %v, want 20", weekly.Percentage)
	}
	if weekly.ResetAt != "2025-03-10T12:00:00Z" {
		t.Errorf("Weekly ResetAt = %v, want 2025-03-10T12:00:00Z", weekly.ResetAt)
	}

	// Verify inferred plan (cap 25 / 100 is unknown)
	if usage.Plan.Name != "Unknown" || usage.Plan.Code != "unknown" {
		t.Errorf("Plan = %+v, want Unknown", usage.Plan)
	}

	// Verify plan inference for GOAT
	rawGOAT := []byte(`{
		"credits": {"monthlyCredits": 100},
		"windowLimits": {
			"fiveHour": {"used": 2, "cap": 14},
			"weekly": {"used": 10, "cap": 35}
		}
	}`)
	usageGOAT, errGOAT := ParseAndFormatUsage(rawGOAT, nil, now)
	if errGOAT != nil {
		t.Fatalf("ParseAndFormatUsage GOAT error: %v", errGOAT)
	}
	if usageGOAT.Plan.Name != "GOAT" || usageGOAT.Plan.Code != "goat" {
		t.Errorf("GOAT Plan = %+v, want name=GOAT code=goat", usageGOAT.Plan)
	}
}

func TestFetchCreditsRaw_FallbackHTTP(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/internal/billing/credits" {
			t.Errorf("unexpected path: %s", r.URL.Path)
			http.NotFound(w, r)
			return
		}

		cookie := r.Header.Get("Cookie")
		expectedCookie := "__Secure-commandcode_prod_.session_token=test-session-123"
		if cookie != expectedCookie {
			t.Errorf("Cookie = %q, want %q", cookie, expectedCookie)
			w.WriteHeader(http.StatusUnauthorized)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"credits":{"monthlyCredits":100},"windowLimits":{"fiveHour":{"used":1,"cap":10}}}`))
	}))
	defer ts.Close()

	// Ensure hostCaller is nil for fallback test
	SetHostCaller(nil)
	SetDefaultHTTPClient(ts.Client())

	body, status, err := FetchCreditsRaw(context.Background(), ts.URL, "test-session-123", "")
	if err != nil {
		t.Fatalf("FetchCreditsRaw error: %v", err)
	}
	if status != http.StatusOK {
		t.Errorf("status = %d, want 200", status)
	}
	if len(body) == 0 {
		t.Fatal("expected non-empty body")
	}

	usage, errParse := ParseAndFormatUsage(body, nil, time.Time{})
	if errParse != nil {
		t.Fatalf("ParseAndFormatUsage error: %v", errParse)
	}
	if usage.Credits.MonthlyCredits != 100 {
		t.Errorf("MonthlyCredits = %v, want 100", usage.Credits.MonthlyCredits)
	}
}

func TestFetchCreditsRaw_MissingToken(t *testing.T) {
	_, status, err := FetchCreditsRaw(context.Background(), "", "", "")
	if err == nil {
		t.Fatal("expected error for missing token")
	}
	if status != http.StatusBadRequest {
		t.Errorf("status = %d, want 400", status)
	}
}

func TestFetchCreditsRaw_HostCaller(t *testing.T) {
	mockResponsePayload := []byte(`{"credits":{"monthlyCredits":500},"windowLimits":{"fiveHour":{"used":5,"cap":50}}}`)

	SetHostCaller(func(method string, payload []byte) ([]byte, error) {
		if method != "host.http.do" {
			t.Errorf("method = %s, want host.http.do", method)
		}
		hostResp := HostHTTPResponse{
			StatusCode: http.StatusOK,
			Body:       mockResponsePayload,
		}
		respJSON, _ := json.Marshal(hostResp)
		return json.Marshal(Envelope{OK: true, Result: respJSON})
	})
	defer SetHostCaller(nil)

	body, status, err := FetchCreditsRaw(context.Background(), "https://api.commandcode.ai", "mock-token", "cb-123")
	if err != nil {
		t.Fatalf("FetchCreditsRaw with hostCaller error: %v", err)
	}
	if status != http.StatusOK {
		t.Errorf("status = %d, want 200", status)
	}
	if string(body) != string(mockResponsePayload) {
		t.Errorf("body = %s, want %s", string(body), string(mockResponsePayload))
	}
}

func TestPlanFromWindowLimits(t *testing.T) {
	trueVal := true
	falseVal := false

	tests := []struct {
		name        string
		fiveHourCap float64
		weeklyCap   float64
		limited     *bool
		wantName    string
		wantCode    string
	}{
		{
			name:        "GOAT plan",
			fiveHourCap: 14,
			weeklyCap:   35,
			limited:     &trueVal,
			wantName:    "GOAT",
			wantCode:    "goat",
		},
		{
			name:        "Pro plan",
			fiveHourCap: 16,
			weeklyCap:   40,
			limited:     nil,
			wantName:    "Pro",
			wantCode:    "pro",
		},
		{
			name:        "Max 20x plan",
			fiveHourCap: 90,
			weeklyCap:   180,
			limited:     &trueVal,
			wantName:    "Max 20×",
			wantCode:    "max_20x",
		},
		{
			name:        "Max 10x plan",
			fiveHourCap: 45,
			weeklyCap:   90,
			limited:     nil,
			wantName:    "Max 10×",
			wantCode:    "max_10x",
		},
		{
			name:        "Go plan",
			fiveHourCap: 3,
			weeklyCap:   6,
			limited:     nil,
			wantName:    "Go",
			wantCode:    "go",
		},
		{
			name:        "Provider pay-as-you-go plan",
			fiveHourCap: 0,
			weeklyCap:   0,
			limited:     &falseVal,
			wantName:    "Provider",
			wantCode:    "provider",
		},
		{
			name:        "Provider plan with caps set but limited=false",
			fiveHourCap: 14,
			weeklyCap:   35,
			limited:     &falseVal,
			wantName:    "Provider",
			wantCode:    "provider",
		},
		{
			name:        "Float tolerance test",
			fiveHourCap: 13.999,
			weeklyCap:   35.001,
			limited:     nil,
			wantName:    "GOAT",
			wantCode:    "goat",
		},
		{
			name:        "Unknown caps",
			fiveHourCap: 10,
			weeklyCap:   20,
			limited:     nil,
			wantName:    "Unknown",
			wantCode:    "unknown",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := PlanFromWindowLimits(tt.fiveHourCap, tt.weeklyCap, tt.limited)
			if got.Name != tt.wantName || got.Code != tt.wantCode {
				t.Errorf("PlanFromWindowLimits(%v, %v, %v) = %+v, want name=%q code=%q",
					tt.fiveHourCap, tt.weeklyCap, tt.limited, got, tt.wantName, tt.wantCode)
			}
		})
	}
}