chore: directory migration - gateway→server, web→client/electron
This commit is contained in:
253
server/internal/repo/agent_disable_test.go
Normal file
253
server/internal/repo/agent_disable_test.go
Normal file
@ -0,0 +1,253 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// 停用是可逆的「归档」，不是删除。这组测试钉住三件事：
|
||||
// 停用后从候选里消失、密钥被撤销、重新注册不能复活它。
|
||||
|
||||
func TestSetAgentDisabledHidesFromCandidates(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for _, n := range []string{"keeper", "goner"} {
|
||||
if err := CreateOrUpdateAgent(ctx, n, "s", "test", nil); err != nil {
|
||||
t.Fatalf("注册 %s: %v", n, err)
|
||||
}
|
||||
}
|
||||
|
||||
if _, err := SetAgentDisabled(ctx, "goner", true); err != nil {
|
||||
t.Fatalf("停用: %v", err)
|
||||
}
|
||||
|
||||
// 默认列表（地址补全、GET /agents、可授权范围都走这条）不含已停用的
|
||||
got, err := ListAgents(ctx, "")
|
||||
if err != nil {
|
||||
t.Fatalf("ListAgents: %v", err)
|
||||
}
|
||||
names := map[string]bool{}
|
||||
for _, a := range got {
|
||||
names[a.Name] = true
|
||||
}
|
||||
if names["goner"] {
|
||||
t.Error("已停用的 Agent 仍出现在默认列表里 —— 人会把任务派给一个不会响应的地址")
|
||||
}
|
||||
if !names["keeper"] {
|
||||
t.Error("停用一个把别的也弄没了")
|
||||
}
|
||||
|
||||
// statusFilter="all" 时要能看到 —— 那是管理页恢复它的唯一入口
|
||||
all, err := ListAgents(ctx, "all")
|
||||
if err != nil {
|
||||
t.Fatalf("ListAgents(all): %v", err)
|
||||
}
|
||||
found := false
|
||||
for _, a := range all {
|
||||
if a.Name == "goner" {
|
||||
found = true
|
||||
if a.Status != "disabled" {
|
||||
t.Errorf("状态应为 disabled，实际 %q", a.Status)
|
||||
}
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Error("statusFilter=all 也看不到已停用的，就再也无法恢复它了")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetAgentDisabledRevokesKeys(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := CreateOrUpdateAgent(ctx, "bot", "s", "test", nil); err != nil {
|
||||
t.Fatalf("注册: %v", err)
|
||||
}
|
||||
admin := seedAdminForTest(t, ctx)
|
||||
for i := 0; i < 2; i++ {
|
||||
if _, err := CreateAgentKey(ctx, "bot", "permanent", "k", 0, admin, ""); err != nil {
|
||||
t.Fatalf("建密钥: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
revoked, err := SetAgentDisabled(ctx, "bot", true)
|
||||
if err != nil {
|
||||
t.Fatalf("停用: %v", err)
|
||||
}
|
||||
if revoked != 2 {
|
||||
t.Errorf("应撤销 2 把密钥，实际 %d", revoked)
|
||||
}
|
||||
|
||||
keys, err := ListAgentKeys(ctx, "bot")
|
||||
if err != nil {
|
||||
t.Fatalf("ListAgentKeys: %v", err)
|
||||
}
|
||||
if len(keys) != 0 {
|
||||
t.Errorf("停用后仍留着 %d 把密钥 —— 插件还能用它调 /mail/send，"+
|
||||
"停用的语义是「不再参与工作」而不只是「不出现在补全里」", len(keys))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDisabledAgentCannotReRegister(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := CreateOrUpdateAgent(ctx, "bot", "s", "test", nil); err != nil {
|
||||
t.Fatalf("首次注册: %v", err)
|
||||
}
|
||||
if _, err := SetAgentDisabled(ctx, "bot", true); err != nil {
|
||||
t.Fatalf("停用: %v", err)
|
||||
}
|
||||
|
||||
// 插件启动时会重新注册。不拒的话 status 被写回 online，停用等于没做。
|
||||
err := CreateOrUpdateAgent(ctx, "bot", "s", "test", nil)
|
||||
if !errors.Is(err, ErrAgentDisabled) {
|
||||
t.Fatalf("已停用的 Agent 重新注册应当被拒，实际 err=%v", err)
|
||||
}
|
||||
|
||||
disabled, err := AgentDisabled(ctx, "bot")
|
||||
if err != nil {
|
||||
t.Fatalf("AgentDisabled: %v", err)
|
||||
}
|
||||
if !disabled {
|
||||
t.Error("注册尝试把停用状态冲掉了")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeartbeatDoesNotReviveDisabledAgent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := CreateOrUpdateAgent(ctx, "bot", "s", "test", nil); err != nil {
|
||||
t.Fatalf("注册: %v", err)
|
||||
}
|
||||
if _, err := SetAgentDisabled(ctx, "bot", true); err != nil {
|
||||
t.Fatalf("停用: %v", err)
|
||||
}
|
||||
|
||||
// 心跳是 30 秒一次的。不排除 disabled 的话停用最多维持半分钟。
|
||||
if _, err := HeartbeatAgent(ctx, "bot"); err != nil {
|
||||
t.Fatalf("心跳本身不该报错: %v", err)
|
||||
}
|
||||
|
||||
disabled, _ := AgentDisabled(ctx, "bot")
|
||||
if !disabled {
|
||||
t.Error("心跳把已停用的 Agent 改回在线了")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestoreAgentGoesOfflineNotOnline(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := CreateOrUpdateAgent(ctx, "bot", "s", "test", nil); err != nil {
|
||||
t.Fatalf("注册: %v", err)
|
||||
}
|
||||
if _, err := SetAgentDisabled(ctx, "bot", true); err != nil {
|
||||
t.Fatalf("停用: %v", err)
|
||||
}
|
||||
if _, err := SetAgentDisabled(ctx, "bot", false); err != nil {
|
||||
t.Fatalf("恢复: %v", err)
|
||||
}
|
||||
|
||||
all, _ := ListAgents(ctx, "all")
|
||||
for _, a := range all {
|
||||
if a.Name != "bot" {
|
||||
continue
|
||||
}
|
||||
// 恢复成 online 会让界面显示一个其实没在跑的 Agent 为在线；
|
||||
// 它是否真的活着由下一次心跳决定。
|
||||
if a.Status != "offline" {
|
||||
t.Errorf("恢复后应为 offline，实际 %q", a.Status)
|
||||
}
|
||||
}
|
||||
|
||||
// 恢复后能重新注册
|
||||
if err := CreateOrUpdateAgent(ctx, "bot", "s", "test", nil); err != nil {
|
||||
t.Errorf("恢复后应当能重新注册: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetAgentDisabledUnknownAgent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := SetAgentDisabled(ctx, "nope", true)
|
||||
if !errors.Is(err, sql.ErrNoRows) {
|
||||
t.Errorf("停用不存在的 Agent 应回 ErrNoRows，实际 %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 停用不得动邮件与会话 —— 往来里有一半是人自己写的。
|
||||
func TestSetAgentDisabledKeepsMailAndSessions(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := CreateOrUpdateAgent(ctx, "bot", "s", "test", nil); err != nil {
|
||||
t.Fatalf("注册: %v", err)
|
||||
}
|
||||
sid, err := CreateSession(ctx, nil, "bot", "一件事", "/tmp/ws")
|
||||
if err != nil {
|
||||
t.Fatalf("建会话: %v", err)
|
||||
}
|
||||
if _, err := CreateMail(ctx, sid, nil,
|
||||
"human", "", "bot", "/tmp/ws", "主题", "正文", nil); err != nil {
|
||||
t.Fatalf("建邮件: %v", err)
|
||||
}
|
||||
|
||||
if _, err := SetAgentDisabled(ctx, "bot", true); err != nil {
|
||||
t.Fatalf("停用: %v", err)
|
||||
}
|
||||
|
||||
mails, err := ListInbox(ctx, "bot", "all", 10)
|
||||
if err != nil {
|
||||
t.Fatalf("ListInbox: %v", err)
|
||||
}
|
||||
if len(mails) != 1 {
|
||||
t.Errorf("停用把邮件删了：剩 %d 封。那些往来里有一半是人自己写的", len(mails))
|
||||
}
|
||||
}
|
||||
|
||||
// 模型范围与平台会话镜像也保留：恢复后不必重配。
|
||||
func TestSetAgentDisabledKeepsModelScope(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := CreateOrUpdateAgent(ctx, "bot", "s", "test", nil); err != nil {
|
||||
t.Fatalf("注册: %v", err)
|
||||
}
|
||||
if err := SetAllowedModels(ctx, "bot", []ModelRef{{Provider: "p", Model: "m"}}); err != nil {
|
||||
t.Fatalf("设范围: %v", err)
|
||||
}
|
||||
|
||||
if _, err := SetAgentDisabled(ctx, "bot", true); err != nil {
|
||||
t.Fatalf("停用: %v", err)
|
||||
}
|
||||
|
||||
allowed, err := ListAllowedModels(ctx, "bot")
|
||||
if err != nil {
|
||||
t.Fatalf("ListAllowedModels: %v", err)
|
||||
}
|
||||
if len(allowed) != 1 {
|
||||
t.Errorf("停用把模型范围清了，恢复后管理员得重配一遍：%+v", allowed)
|
||||
}
|
||||
}
|
||||
|
||||
// seedAdminForTest 插一个管理员并返回它的 user_id（CreateAgentKey 要 created_by）。
|
||||
func seedAdminForTest(t *testing.T, ctx context.Context) uuid.UUID {
|
||||
t.Helper()
|
||||
var id uuid.UUID
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`INSERT INTO users (username, display_name, password_hash, role)
|
||||
VALUES ('key-admin', 'Admin', 'x', 'admin') RETURNING user_id`).Scan(&id)
|
||||
if err != nil {
|
||||
t.Fatalf("seed admin: %v", err)
|
||||
}
|
||||
return id
|
||||
}
|
||||
Reference in New Issue
Block a user