Vectrum
Sign In
185 lines · 6.3 KB · Go
package main

import (
	"bytes"
	"net/http"
	"net/http/httptest"
	"regexp"
	"slices"
	"strings"
	"testing"

	"github.com/johannSo/codey/internal/llm"
	"github.com/johannSo/codey/internal/serve"
)

// t3VersionPattern is the regular expression T3 Code parses `--version` with.
// It is copied here rather than described, because a version line that does
// not match it reads as a binary that is installed and broken.
var t3VersionPattern = regexp.MustCompile(`\bv?(\d+\.\d+\.\d+)\b`)

func TestVersionMatchesWhatTheClientLooksFor(t *testing.T) {
	var out bytes.Buffer
	printVersion(&out)

	m := t3VersionPattern.FindStringSubmatch(out.String())
	if m == nil {
		t.Fatalf("no version in %q", out.String())
	}
	if !strings.Contains(out.String(), "codey") {
		t.Errorf("the line should say whose version it is: %q", out.String())
	}
}

func TestScanArgsReadsEveryShapeTheClientSends(t *testing.T) {
	// The five launches T3 Code's Grok driver produces, one per runtime mode.
	cases := []struct {
		name string
		args []string
		want serve.PermissionMode
	}{
		{"default", []string{"agent", "stdio"}, serve.ModeAsk},
		{"approval required", []string{"--permission-mode", "default", "agent", "stdio"}, serve.ModeAsk},
		{"accept edits", []string{"--permission-mode", "acceptEdits", "agent", "stdio"}, serve.ModeAcceptEdits},
		{"auto", []string{"--permission-mode", "auto", "agent", "stdio"}, serve.ModeAuto},
		{"full access", []string{"agent", "--always-approve", "stdio"}, serve.ModeFullAccess},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			inv := scanArgs(tc.args)
			if inv.Kind != invAgent {
				t.Fatalf("kind = %v, want an agent run", inv.Kind)
			}
			if inv.Permission != tc.want {
				t.Errorf("permission = %q, want %q", inv.Permission, tc.want)
			}
			if len(inv.Rest) != 0 {
				t.Errorf("leftover args %v — flag parsing would choke on them", inv.Rest)
			}
		})
	}
}

func TestScanArgsKeepsOurOwnFlagsFromTheLaunchArguments(t *testing.T) {
	// The launch arguments are a free-text box in a settings panel: whatever
	// the user typed arrives mixed in with the subcommand.
	inv := scanArgs([]string{"--permission-mode", "auto", "agent", "stdio", "-effort", "high", "-no-web"})

	if inv.Kind != invAgent || inv.Permission != serve.ModeAuto {
		t.Fatalf("invocation = %+v", inv)
	}
	want := []string{"-effort", "high", "-no-web"}
	if !slices.Equal(inv.Rest, want) {
		t.Errorf("rest = %v, want %v", inv.Rest, want)
	}
}

func TestScanArgsTakesThePermissionModeJoinedByAnEquals(t *testing.T) {
	if got := scanArgs([]string{"--permission-mode=acceptEdits", "agent", "stdio"}); got.Permission != serve.ModeAcceptEdits {
		t.Errorf("permission = %q", got.Permission)
	}
}

func TestAnUnknownPermissionModeIsTheStrictOne(t *testing.T) {
	// A mode we have never heard of is not a reason to stop asking.
	if got := scanArgs([]string{"--permission-mode", "yolo-supreme", "agent", "stdio"}); got.Permission != serve.ModeAsk {
		t.Errorf("permission = %q, want %q", got.Permission, serve.ModeAsk)
	}
}

func TestTheProbesAreRecognisedOnlyOnTheirOwn(t *testing.T) {
	if got := scanArgs([]string{"--version"}).Kind; got != invVersion {
		t.Errorf("--version kind = %v", got)
	}
	if got := scanArgs([]string{"models"}).Kind; got != invModels {
		t.Errorf("models kind = %v", got)
	}
	// A directory that happens to be called models is a directory.
	inv := scanArgs([]string{"-C", "models"})
	if inv.Kind != invNormal {
		t.Errorf("kind = %v, want a normal run", inv.Kind)
	}
	if !slices.Equal(inv.Rest, []string{"-C", "models"}) {
		t.Errorf("a normal run must get its line back untouched: %v", inv.Rest)
	}
}

func TestHalfAnAgentInvocationIsNotOne(t *testing.T) {
	// "agent" without "stdio" is somebody's file name, not a subcommand.
	for _, args := range [][]string{{"agent"}, {"stdio"}, {"-p", "explain the agent"}} {
		if got := scanArgs(args).Kind; got != invNormal {
			t.Errorf("scanArgs(%v) = %v, want a normal run", args, got)
		}
	}
}

// testClient has a credential and an endpoint that refuses to list anything,
// which is the interesting case: the probe still has to answer.
func testClient(t *testing.T) *llm.Client {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
	}))
	t.Cleanup(srv.Close)
	return llm.New(llm.WithBaseURL(srv.URL), llm.WithAPIKey("k"), llm.WithModel("codey-test-model"))
}

func noCredentialClient(t *testing.T) *llm.Client {
	t.Helper()
	c := llm.New(llm.WithBaseURL(llm.DefaultBaseURL))
	c.ClearCredentialsForTest()
	return c
}

// t3ModelBullet is how T3 Code reads a line of the model listing, and
// t3LoggedIn is how it reads the login state off the same output.
var (
	t3ModelBullet = regexp.MustCompile(`^\s*[*-]\s+(\S+)(.*)$`)
	t3LoggedIn    = regexp.MustCompile(`(?i)you are logged in`)
	t3LoggedOut   = regexp.MustCompile(`(?i)not authenticated|not logged in`)
)

// parseLikeT3 is a transliteration of the client's parser, so the test fails
// for the same reason the provider card would.
func parseLikeT3(out string) (loggedIn bool, models []string, defaultModel string) {
	loggedIn = t3LoggedIn.MatchString(out) && !t3LoggedOut.MatchString(out)
	for _, line := range strings.Split(out, "\n") {
		m := t3ModelBullet.FindStringSubmatch(line)
		if m == nil {
			continue
		}
		models = append(models, m[1])
		if strings.Contains(strings.ToLower(m[2]), "(default)") {
			defaultModel = m[1]
		}
	}
	return
}

func TestModelsOutputParsesTheWayTheClientParsesIt(t *testing.T) {
	// No endpoint is reachable here, so this exercises the fallback: the one
	// model we were configured with, and a login state that is still true.
	var out bytes.Buffer
	printModels(&out, testClient(t))

	loggedIn, models, def := parseLikeT3(out.String())
	if !loggedIn {
		t.Errorf("the client would read this as logged out:\n%s", out.String())
	}
	if len(models) != 1 || def != models[0] {
		t.Errorf("models = %v, default = %q, from:\n%s", models, def, out.String())
	}
}

func TestLoggedOutSaysSoRatherThanFailing(t *testing.T) {
	var out bytes.Buffer
	printModels(&out, noCredentialClient(t))

	if loggedIn, _, _ := parseLikeT3(out.String()); loggedIn {
		t.Errorf("a client with no credential reported itself logged in:\n%s", out.String())
	}
	if !t3LoggedOut.MatchString(out.String()) {
		t.Errorf("the client would read this as unknown rather than logged out:\n%s", out.String())
	}
}