pulse/internal/ai/tools/control_test.go

package tools

import (
	"context"
	"encoding/json"
	"strings"
	"testing"
	"time"

	"github.com/rcourtman/pulse-go-rewrite/internal/agentexec"
	"github.com/rcourtman/pulse-go-rewrite/internal/models"
	"github.com/rcourtman/pulse-go-rewrite/internal/unifiedresources"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
)

func mustParseJSONMapControl(t *testing.T, text string) map[string]interface{} {
	t.Helper()
	var out map[string]interface{}
	assert.NoError(t, json.Unmarshal([]byte(text), &out))
	return out
}

func TestPulseToolExecutor_ExecuteListBackups(t *testing.T) {
	backupProv := &mockBackupProvider{}
	exec := NewPulseToolExecutor(ExecutorConfig{
		BackupProvider: backupProv,
		ControlLevel:   ControlLevelReadOnly,
	})

	expectedBackups := models.Backups{
		PBS: []models.PBSBackup{
			{VMID: "101", BackupType: "vm", BackupTime: time.Now(), Instance: "pbs1", Datastore: "ds1", Size: 1024 * 1024 * 1024},
		},
		PVE: models.PVEBackups{
			StorageBackups: []models.StorageBackup{
				{VMID: 102, Time: time.Now(), Size: 2 * 1024 * 1024 * 1024, Storage: "local"},
			},
			BackupTasks: []models.BackupTask{
				{VMID: 101, Node: "node1", Status: "OK", StartTime: time.Now()},
			},
		},
	}
	backupProv.On("GetBackups").Return(expectedBackups)
	backupProv.On("GetPBSInstances").Return([]models.PBSInstance{})

	// Use pulse_storage tool with type: "backups"
	result, err := exec.ExecuteTool(context.Background(), "pulse_storage", map[string]interface{}{
		"type": "backups",
	})
	assert.NoError(t, err)
	assert.False(t, result.IsError)
}

func TestPulseToolExecutor_ExecuteControlGuest(t *testing.T) {
	stateProv := &mockStateProvider{}
	agentSrv := &mockAgentServer{}
	exec := NewPulseToolExecutor(ExecutorConfig{
		StateProvider:    stateProv,
		AgentServer:      agentSrv,
		ControlLevel:     ControlLevelAutonomous,
		ActionAuditStore: unifiedresources.NewMemoryStore(),
	})

	state := models.StateSnapshot{
		VMs: []models.VM{
			{VMID: 100, Name: "test-vm", Node: "node1", Status: "running", Instance: "pve1"},
		},
	}
	stateProv.On("ReadSnapshot").Return(state)

	agentSrv.On("GetConnectedAgents").Return([]agentexec.ConnectedAgent{
		{AgentID: "agent1", Hostname: "node1"},
	})

	agentSrv.On("ExecuteCommand", mock.Anything, "agent1", mock.MatchedBy(func(payload agentexec.ExecuteCommandPayload) bool {
		return payload.Command == "qm stop 100"
	})).Return(&agentexec.CommandResultPayload{
		Stdout:   "OK",
		ExitCode: 0,
	}, nil)

	agentSrv.On("ExecuteCommand", mock.Anything, "agent1", mock.MatchedBy(func(payload agentexec.ExecuteCommandPayload) bool {
		return payload.Command == "qm status 100"
	})).Return(&agentexec.CommandResultPayload{
		Stdout:   "status: stopped",
		ExitCode: 0,
	}, nil)

	// Use pulse_control tool with type: "guest"
	result, err := exec.ExecuteTool(context.Background(), "pulse_control", map[string]interface{}{
		"type":     "guest",
		"guest_id": "100",
		"action":   "stop",
	})
	assert.NoError(t, err)
	resp := mustParseJSONMapControl(t, result.Content[0].Text)
	assert.Equal(t, true, resp["success"])
	if v, ok := resp["verification"].(map[string]interface{}); ok {
		assert.Equal(t, true, v["ok"])
	}
}

func TestPulseToolExecutor_ExecuteControlDocker(t *testing.T) {
	stateProv := &mockStateProvider{}
	agentSrv := &mockAgentServer{}
	exec := NewPulseToolExecutor(ExecutorConfig{
		StateProvider:    stateProv,
		AgentServer:      agentSrv,
		ControlLevel:     ControlLevelAutonomous,
		ActionAuditStore: unifiedresources.NewMemoryStore(),
	})

	state := models.StateSnapshot{
		DockerHosts: []models.DockerHost{
			{
				ID:       "host1",
				Hostname: "docker1",
				Containers: []models.DockerContainer{
					{ID: "c1", Name: "nginx"},
				},
			},
		},
	}
	stateProv.On("ReadSnapshot").Return(state)

	agentSrv.On("GetConnectedAgents").Return([]agentexec.ConnectedAgent{
		{AgentID: "agent1", Hostname: "docker1"},
	})

	agentSrv.On("ExecuteCommand", mock.Anything, "agent1", mock.MatchedBy(func(payload agentexec.ExecuteCommandPayload) bool {
		return payload.Command == "docker restart 'nginx'"
	})).Return(&agentexec.CommandResultPayload{
		Stdout:   "OK",
		ExitCode: 0,
	}, nil)

	agentSrv.On("ExecuteCommand", mock.Anything, "agent1", mock.MatchedBy(func(payload agentexec.ExecuteCommandPayload) bool {
		return strings.HasPrefix(payload.Command, "docker inspect -f")
	})).Return(&agentexec.CommandResultPayload{
		Stdout:   "running true",
		ExitCode: 0,
	}, nil)

	// Use pulse_docker tool with action: "control"
	result, err := exec.ExecuteTool(context.Background(), "pulse_docker", map[string]interface{}{
		"action":    "control",
		"container": "nginx",
		"operation": "restart",
	})
	assert.NoError(t, err)
	resp := mustParseJSONMapControl(t, result.Content[0].Text)
	assert.Equal(t, true, resp["success"])
	if v, ok := resp["verification"].(map[string]interface{}); ok {
		assert.Equal(t, true, v["ok"])
	}
}