Files
ai-agent/internal/services/agent_revision_service_test.go
T
mlogclub 34051a4631 feat: Enhance AI Agent and Channel Management
- Updated labels in the AI Agents dashboard for clarity, changing "流程状态" to "Playbook 状态" and "未发布流程" to "未发布 Playbook".
- Introduced AI Agent rollout percentage management in channel editing, allowing users to set and rollback rollout percentages.
- Added new API endpoints for rolling back AI Agent rollout and fetching agent run metrics.
- Implemented new UI components for displaying agent run details, including status, duration, and input/output tokens.
- Enhanced type definitions for AdminChannel and AIAgent to include rollout percentages and runtime modes.
- Updated navigation to include a section for agent runs.
- Added new translations for agent run features in both English and Chinese.
2026-07-25 12:04:06 +08:00

51 lines
2.1 KiB
Go

package services
import (
"encoding/json"
"strings"
"testing"
"agent-desk/internal/models"
"agent-desk/internal/pkg/enums"
"github.com/glebarez/sqlite"
"github.com/mlogclub/simple/sqls"
"gorm.io/gorm"
)
func TestAgentRevisionServiceRestoresPublishedSnapshotAndKeepsAPIKey(t *testing.T) {
db, err := gorm.Open(sqlite.Open("file:"+strings.ReplaceAll(t.Name(), "/", "_")+"?mode=memory&cache=shared"), &gorm.Config{})
if err != nil {
t.Fatalf("open sqlite: %v", err)
}
if err := db.AutoMigrate(&models.AgentRevision{}); err != nil {
t.Fatalf("auto migrate: %v", err)
}
sqls.SetDB(db)
definition := agentRevisionDefinition{
Agent: agentRevisionAgent{
Name: "published agent", AIConfigID: 8, RuntimeMode: string(enums.AIAgentRuntimeModeAutonomous),
MaxSteps: 5, ContextWindow: 9, SystemPrompt: "published instruction", KnowledgeIDs: "4", ReplyTimeoutSeconds: 90,
},
Model: agentRevisionModel{ConfigID: 8, Provider: string(enums.AIProviderOpenAI), BaseURL: "https://published.example/v1", ModelType: string(enums.AIModelTypeLLM), ModelName: "published-model", TimeoutMS: 12000},
}
data, err := json.Marshal(definition)
if err != nil {
t.Fatalf("marshal definition: %v", err)
}
revision := &models.AgentRevision{AgentID: 7, Revision: 1, Status: enums.StatusOk, Definition: string(data)}
if err := db.Create(revision).Error; err != nil {
t.Fatalf("create revision: %v", err)
}
snapshot, err := AgentRevisionService.ResolvePublishedSnapshot(models.AIAgent{ID: 7, PublishedRevisionID: revision.ID, SystemPrompt: "draft instruction"}, models.AIConfig{ID: 8, APIKey: "rotated-secret", ModelName: "draft-model"})
if err != nil {
t.Fatalf("ResolvePublishedSnapshot: %v", err)
}
if snapshot.Agent.SystemPrompt != "published instruction" || snapshot.Agent.MaxSteps != 5 || snapshot.Agent.ReplyTimeoutSeconds != 90 {
t.Fatalf("agent snapshot not restored: %#v", snapshot.Agent)
}
if snapshot.AIConfig.ModelName != "published-model" || snapshot.AIConfig.BaseURL != "https://published.example/v1" || snapshot.AIConfig.APIKey != "rotated-secret" {
t.Fatalf("model snapshot not restored safely: %#v", snapshot.AIConfig)
}
}