magpie/internal/provider/commandcode_go_test.go

package provider

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

// A Command Code account on the Go plan is asked at /alpha/generate, with
// its key and Go's models; the plan is read from billing/subscriptions
// and kept, the one its sign-in said standing in when it can't be read.
// Every other plan stays on the Provider API.
func TestCommandCodeGoPlan(t *testing.T) {
	home := signIn(t)
	writeFile(t, filepath.Join(home, ".commandcode", "auth.json"), map[string]any{"apiKey": "go-key", "userName": "gouser"})
	var asked atomic.Int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		key := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ")
		ok := func(v any) { _ = json.NewEncoder(w).Encode(v) }
		switch {
		case r.URL.Path == "/provider/v1/models":
			t.Errorf("the Provider API asked for %q's models", key)
			w.WriteHeader(403)
		case r.URL.Path != "/alpha/billing/subscriptions":
			w.WriteHeader(404)
		case key == "go-key":
			asked.Add(1)
			ok(map[string]any{"success": true, "data": map[string]any{"planId": "individual-go-monthly", "status": "active"}})
		case key == "goat-key":
			ok(map[string]any{"success": true, "data": map[string]any{"planId": "individual-goat-monthly", "status": "active"}})
		case key == "pro-key":
			ok(map[string]any{"success": true, "data": map[string]any{"planId": "individual-pro-monthly", "status": "active"}})
		default:
			w.WriteHeader(503)
		}
	}))
	defer srv.Close()
	oldAPI := cmdAPI
	cmdAPI = srv.URL
	defer func() { cmdAPI = oldAPI }()
	cmdPlansSeen.Lock()
	cmdPlansSeen.m = map[string]cmdSeen{}
	cmdPlansSeen.Unlock()
	ctx := context.Background()

	// the CLI's own account, with no plan saved, is asked which it is on
	p, found := find(All(), CommandCodePlanID)
	if !found {
		t.Fatal("no Command Code account")
	}
	api, key, ok := CommandCodeGenerate(ctx, p)
	if !ok || api != srv.URL || key != "go-key" {
		t.Fatalf("go: %v %q %q", ok, api, key)
	}
	ms, err := p.Account.fetch(ctx)
	if err != nil || len(ms) != len(cmdGoModels) || ms[0].ID != "gpt-6-luna" {
		t.Fatalf("go models: %v %+v", err, ms)
	}
	if got := p.Account.models(); len(got) != len(cmdGoModels) {
		t.Fatalf("go models before a fetch: %+v", got)
	}
	// read once, then kept
	if asked.Load() != 1 {
		t.Fatalf("billing/subscriptions asked %d times", asked.Load())
	}
	if p, _ := find(All(), CommandCodePlanID); p.Account.Plan != "Go" {
		t.Fatalf("plan: %q", p.Account.Plan)
	}
	// and asked again once it is old
	cmdPlansSeen.Lock()
	cmdPlansSeen.m["go-key"] = cmdSeen{"Go", time.Now().Add(-cmdPlanKeep)}
	cmdPlansSeen.Unlock()
	if _, _, ok := CommandCodeGenerate(ctx, p); !ok || asked.Load() != 2 {
		t.Fatalf("again: %v, asked %d", ok, asked.Load())
	}

	for _, c := range []struct {
		key, saved string
		want       bool
	}{
		{"pro-key", "", false},
		{"goat-key", "", false},  // GOAT is not Go
		{"pro-key", "Go", false}, // what billing says, over the sign-in
		{"down-key", "Go", true}, // can't be read: the sign-in's plan
		{"down-key", "Max", false},
	} {
		cmdPlansSeen.Lock()
		delete(cmdPlansSeen.m, c.key)
		cmdPlansSeen.Unlock()
		p := cmdProvider("u", c.saved, cmdAuth{APIKey: c.key})
		_, key, ok := CommandCodeGenerate(ctx, p)
		if ok != c.want || (ok && key != c.key) {
			t.Errorf("%s saved %q: %v %q", c.key, c.saved, ok, key)
		}
		if !ok && p.Anthropic != srv.URL+"/provider" {
			t.Errorf("%s: not on the Provider API: %+v", c.key, p)
		}
	}
	// a keyed provider is never asked there
	if _, _, ok := CommandCodeGenerate(ctx, Provider{ID: "commandcode", Key: "k"}); ok {
		t.Fatal("keyed provider")
	}
}

// Reading the usage keeps the plan it read, so the account is asked where
// its plan is without another request.
func TestCommandCodeQuotaKeepsPlan(t *testing.T) {
	signIn(t)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/alpha/billing/subscriptions":
			_ = json.NewEncoder(w).Encode(map[string]any{"data": map[string]any{"planId": "individual-go", "status": "active"}})
		case "/alpha/billing/credits":
			_ = json.NewEncoder(w).Encode(map[string]any{"credits": map[string]any{"monthlyCredits": 4}})
		default:
			w.WriteHeader(404)
		}
	}))
	defer srv.Close()
	oldAPI := cmdAPI
	cmdAPI = srv.URL
	defer func() { cmdAPI = oldAPI }()
	cmdPlansSeen.Lock()
	cmdPlansSeen.m = map[string]cmdSeen{}
	cmdPlansSeen.Unlock()
	q := cmdQuota(context.Background(), Login{User: "u"}, cmdAuth{APIKey: "q-key"})
	if q.Plan != "Go" || len(q.Windows) != 1 || q.Windows[0].Display != "$6.00 / $10.00" {
		t.Fatalf("quota: %+v", q)
	}
	if got := cmdPlanKnown(cmdAuth{APIKey: "q-key"}, ""); got != "Go" {
		t.Fatalf("kept: %q", got)
	}
}