package database import ( "path/filepath" "testing" "time" "cyberstrike-ai/internal/mcp" "go.uber.org/zap" ) func TestCancelOrphanedRunningToolExecutions(t *testing.T) { dbPath := filepath.Join(t.TempDir(), "monitor.db") db, err := NewDB(dbPath, zap.NewNop()) if err != nil { t.Fatalf("NewDB: %v", err) } defer db.Close() start := time.Now().Add(-2 * time.Hour) exec := &mcp.ToolExecution{ ID: "orphan-hydra", ToolName: "hydra", Arguments: map[string]interface{}{"target": "127.0.0.1"}, Status: "running", StartTime: start, } if err := db.SaveToolExecution(exec); err != nil { t.Fatalf("SaveToolExecution: %v", err) } end := time.Now() n, err := db.CancelOrphanedRunningToolExecutions(end, "执行已中断(服务重启)") if err != nil { t.Fatalf("CancelOrphanedRunningToolExecutions: %v", err) } if n != 1 { t.Fatalf("expected 1 row updated, got %d", n) } got, err := db.GetToolExecution("orphan-hydra") if err != nil { t.Fatalf("GetToolExecution: %v", err) } if got.Status != "cancelled" { t.Fatalf("expected cancelled, got %s", got.Status) } if got.EndTime == nil { t.Fatal("expected end_time to be set") } if got.Duration <= 0 { t.Fatalf("expected positive duration, got %v", got.Duration) } } func TestFinalizeStaleRunningToolExecutions_skipsActive(t *testing.T) { dbPath := filepath.Join(t.TempDir(), "monitor.db") db, err := NewDB(dbPath, zap.NewNop()) if err != nil { t.Fatalf("NewDB: %v", err) } defer db.Close() now := time.Now() oldStart := now.Add(-5 * time.Minute) if err := db.SaveToolExecution(&mcp.ToolExecution{ ID: "stale", ToolName: "hydra", Status: "running", StartTime: oldStart, }); err != nil { t.Fatalf("SaveToolExecution stale: %v", err) } if err := db.SaveToolExecution(&mcp.ToolExecution{ ID: "active", ToolName: "hydra", Status: "running", StartTime: oldStart, }); err != nil { t.Fatalf("SaveToolExecution active: %v", err) } active := map[string]struct{}{"active": {}} n, err := db.FinalizeStaleRunningToolExecutions(now, time.Minute, active, "执行已中断(会话已结束)") if err != nil { t.Fatalf("FinalizeStaleRunningToolExecutions: %v", err) } if n != 1 { t.Fatalf("expected 1 stale row updated, got %d", n) } stale, err := db.GetToolExecution("stale") if err != nil { t.Fatalf("GetToolExecution stale: %v", err) } if stale.Status != "cancelled" { t.Fatalf("stale expected cancelled, got %s", stale.Status) } activeExec, err := db.GetToolExecution("active") if err != nil { t.Fatalf("GetToolExecution active: %v", err) } if activeExec.Status != "running" { t.Fatalf("active expected running, got %s", activeExec.Status) } }