Vectrum
Sign In
100 lines · 3.0 KB · Go
package approval

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

// fixedRules rules the same way on every call to one tool.
type fixedRules struct {
	tool    string
	verdict Verdict
}

func (f fixedRules) Check(tool string, _ json.RawMessage) Ruling {
	if tool != f.tool {
		return Ruling{}
	}
	return Ruling{Verdict: f.verdict, Rule: "Rule(x)", Source: "settings.json"}
}

func (f fixedRules) Lines() []string { return []string{"settings.json", "  rule"} }

func TestStandingRules(t *testing.T) {
	tests := []struct {
		name       string
		tool       string
		verdict    Verdict
		yolo       bool
		remembered bool
		wantErr    bool
		wantPrompt int
	}{
		{name: "deny holds under yolo", tool: "bash", verdict: VerdictDeny, yolo: true, wantErr: true},
		{name: "deny holds for a read-only tool", tool: "read_file", verdict: VerdictDeny, wantErr: true},
		{name: "deny beats session memory", tool: "bash", verdict: VerdictDeny, remembered: true, wantErr: true},
		{name: "allow skips the prompt", tool: "bash", verdict: VerdictAllow},
		{name: "ask prompts under yolo", tool: "bash", verdict: VerdictAsk, yolo: true, wantPrompt: 1},
		{name: "ask prompts for a read-only tool", tool: "read_file", verdict: VerdictAsk, wantPrompt: 1},
		{name: "ask beats session memory", tool: "bash", verdict: VerdictAsk, remembered: true, wantPrompt: 1},
		{name: "no verdict prompts as usual", tool: "bash", verdict: NoVerdict, wantPrompt: 1},
		{name: "no verdict still honours memory", tool: "bash", verdict: NoVerdict, remembered: true},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			var asked int
			var forced string
			p := New(func(_ context.Context, r Request) (Decision, error) {
				asked++
				forced = r.Forced
				return Once, nil
			})
			p.YOLO = tc.yolo
			p.Rules = fixedRules{tool: tc.tool, verdict: tc.verdict}
			if tc.remembered {
				p.tools[tc.tool] = true
			}
			err := approve(t, p, tc.tool, `{"command":"ls","path":"a"}`)
			if (err != nil) != tc.wantErr {
				t.Fatalf("err = %v, want error %v", err, tc.wantErr)
			}
			if err != nil && !strings.Contains(err.Error(), "Rule(x) in settings.json") {
				t.Errorf("denial %q should name the rule and its file", err)
			}
			if asked != tc.wantPrompt {
				t.Errorf("prompted %d times, want %d", asked, tc.wantPrompt)
			}
			if tc.verdict == VerdictAsk && forced == "" {
				t.Error("a prompt an ask rule forced should say so")
			}
		})
	}
}

// Headless, an ask rule is a refusal with nobody to put the question to —
// which is what the rule asked for.
func TestAskRuleHeadlessDenies(t *testing.T) {
	p := New(nil)
	p.YOLO = true
	p.Rules = fixedRules{tool: "bash", verdict: VerdictAsk}
	if err := approve(t, p, "bash", `{"command":"ls"}`); err == nil {
		t.Error("an ask rule with no prompt should not run the call")
	}
}

func TestRuleLines(t *testing.T) {
	var p *Policy
	if p.RuleLines() != nil {
		t.Error("nil policy")
	}
	p = New(nil)
	if p.RuleLines() != nil {
		t.Error("no rules")
	}
	p.Rules = fixedRules{}
	if len(p.RuleLines()) != 2 {
		t.Error("lines not passed through")
	}
}