CyberStrikeAI/internal/database/monitor_reconcile_test.go

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)
	}
}