Files
ARTEX-chixuweihu/agent/prompt_test.go
T
Autumn-27andClaude Opus 4.8 0ef8561411 chore: ARTEX 初始化
ARTEX — AI 自主渗透测试系统(Go 后端 + Next.js 前端,单二进制内嵌)。

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
2026-07-26 04:27:45 -04:00

68 lines
2.6 KiB
Go

package agent
import (
"strings"
"testing"
)
func TestRenderSystemOverrideAndFallback(t *testing.T) {
t.Cleanup(func() { PromptOverride = nil })
// no override → built-in default
PromptOverride = nil
if got := renderSystem("planner", "DEFAULT", PlannerVars{Goal: "g"}); got != "DEFAULT" {
t.Fatalf("no override should give default, got %q", got)
}
// override → rendered with vars
PromptOverride = func(k string) (string, bool) {
if k == "planner" {
return "目标:{{.Goal}} 范围:{{.Scope}}", true
}
return "", false
}
if got := renderSystem("planner", "DEFAULT", PlannerVars{Goal: "拿下X", Scope: "*.x.com"}); got != "目标:拿下X 范围:*.x.com" {
t.Fatalf("override render: %q", got)
}
// override referencing a non-catalog var → execution error → fallback to default
PromptOverride = func(k string) (string, bool) { return "{{.NotInCatalog}}", true }
if got := renderSystem("planner", "DEFAULT", PlannerVars{Goal: "x"}); got != "DEFAULT" {
t.Fatalf("bad var should fall back to default, got %q", got)
}
// full plannerSystem path: DB body [A] is honored, then the code-owned tail
// [C] (中间产物输出规约) is ALWAYS appended — editing the body can't drop it.
PromptOverride = func(k string) (string, bool) { return "PLANNER {{.Goal}}", true }
got := plannerSystem("拿下X", "/data")
if !strings.HasPrefix(got, "PLANNER 拿下X") {
t.Fatalf("plannerSystem body not honored: %q", got)
}
if !strings.Contains(got, "中间产物输出规约") || !strings.Contains(got, "/data") {
t.Fatalf("plannerSystem missing code-owned artifact tail: %q", got)
}
// worker dual-text via {{if .ProxyAddr}} in a user template, plus the code tail:
// [B] trafficTool present only with a proxy, [C] artifact spec always present.
PromptOverride = func(k string) (string, bool) {
return "{{if .ProxyAddr}}走代理 {{.ProxyAddr}}{{else}}手动{{end}}", true
}
withProxy := workerSystem("127.0.0.1:8080", "/data")
if !strings.HasPrefix(withProxy, "走代理 127.0.0.1:8080") {
t.Fatalf("worker proxy branch body: %q", withProxy)
}
if !strings.Contains(withProxy, "traffic_search") {
t.Fatalf("worker with proxy should inject trafficTool: %q", withProxy)
}
if !strings.Contains(withProxy, "中间产物输出规约") {
t.Fatalf("worker missing artifact tail: %q", withProxy)
}
noProxy := workerSystem("", "/data")
if !strings.HasPrefix(noProxy, "手动") {
t.Fatalf("worker no-proxy branch body: %q", noProxy)
}
if strings.Contains(noProxy, "traffic_search") {
t.Fatalf("worker without proxy must NOT inject trafficTool: %q", noProxy)
}
}