Vectrum
Sign In
137 lines · 4.2 KB · Go
package approval

import (
	"context"
	"encoding/json"
	"errors"
	"strings"
	"testing"

	"github.com/johannSo/codey/internal/hooks"
)

// recordHooks attaches a recorder and returns the payloads it will fill.
func recordHooks(p *Policy) *[]hooks.Payload {
	var got []hooks.Payload
	p.Hooks = hooks.New("s1")
	p.Hooks.Add(hooks.HandlerFunc(func(_ context.Context, pl hooks.Payload) error {
		got = append(got, pl)
		return nil
	}))
	return &got
}

func TestPromptIsBracketedByBlockedAndUnblocked(t *testing.T) {
	tests := []struct {
		name      string
		decision  Decision
		promptErr error
	}{
		{name: "allowed once", decision: Once},
		{name: "denied", decision: Deny},
		{name: "prompt failed", promptErr: errors.New("turn cancelled")},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			// The prompt writes into the same sequence as the hooks, so the
			// test can tell "Blocked before the human was asked" from
			// "Blocked at some point".
			var order []string
			p := New(func(context.Context, Request) (Decision, error) {
				order = append(order, "prompt")
				return tc.decision, tc.promptErr
			})
			var payloads []hooks.Payload
			p.Hooks = hooks.New("s1")
			p.Hooks.Add(hooks.HandlerFunc(func(_ context.Context, pl hooks.Payload) error {
				order = append(order, string(pl.Event))
				payloads = append(payloads, pl)
				return nil
			}))

			err := p.ApproveAs(context.Background(), "a3", "bash", json.RawMessage(`{"command":"go test ./..."}`))
			if tc.promptErr != nil && !errors.Is(err, tc.promptErr) {
				t.Fatalf("ApproveAs = %v, want the prompt's own error", err)
			}
			if want := "blocked prompt unblocked"; strings.Join(order, " ") != want {
				t.Fatalf("order = %v, want %q", order, want)
			}
			for _, pl := range payloads {
				if pl.Reason != hooks.ReasonApproval {
					t.Errorf("%s.Reason = %q, want %q", pl.Event, pl.Reason, hooks.ReasonApproval)
				}
				if pl.Tool != "bash" {
					t.Errorf("%s.Tool = %q, want bash", pl.Event, pl.Tool)
				}
				if pl.Agent != "a3" {
					t.Errorf("%s.Agent = %q, want the asking subagent", pl.Event, pl.Agent)
				}
			}
			// The command is what the prompt itself shows, and what lets a
			// decision be made from a phone.
			if msg := payloads[0].Message; !strings.Contains(msg, "go test ./...") {
				t.Errorf("Blocked.Message = %q, want the command in it", msg)
			}
		})
	}
}

// Blocked means a human is being waited on. A call nobody was asked about
// must not say so, or a sidebar flashes "blocked" on every read.
func TestNothingFiresWhenNoPromptIsNeeded(t *testing.T) {
	tests := []struct {
		name  string
		setup func(p *Policy)
		tool  string
		input string
	}{
		{name: "read-only tool", tool: "read_file", input: `{"path":"a.go"}`},
		{name: "YOLO", setup: func(p *Policy) { p.YOLO = true }, tool: "bash", input: `{"command":"rm -rf build"}`},
		{
			name: "tool remembered as always allowed",
			setup: func(p *Policy) {
				// The remembering itself goes through the prompt, so its
				// events are drained before the call under test.
				approve(t, p, "bash", `{"command":"ls"}`)
			},
			tool: "bash", input: `{"command":"go build ./..."}`,
		},
		{
			name: "command remembered as always allowed",
			setup: func(p *Policy) {
				p.Prompt = func(context.Context, Request) (Decision, error) { return AlwaysThisCommand, nil }
				approve(t, p, "bash", `{"command":"go test ./..."}`)
			},
			tool: "bash", input: `{"command":"go test ./..."}`,
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			p := New(func(context.Context, Request) (Decision, error) { return AlwaysThisTool, nil })
			got := recordHooks(p)
			if tc.setup != nil {
				tc.setup(p)
			}
			*got = (*got)[:0]

			if err := approve(t, p, tc.tool, tc.input); err != nil {
				t.Fatalf("%s was denied: %v", tc.tool, err)
			}
			if len(*got) != 0 {
				t.Errorf("fired %+v for a call nobody was asked about", *got)
			}
		})
	}
}

// A headless denial is instant, so there is no wait to announce.
func TestNilPromptFiresNothing(t *testing.T) {
	p := New(nil)
	got := recordHooks(p)
	if err := approve(t, p, "bash", `{"command":"ls"}`); err == nil {
		t.Fatal("a mutating call was allowed with no way to ask")
	}
	if len(*got) != 0 {
		t.Errorf("fired %+v with no prompt installed", *got)
	}
}