feat: Tenable-style patch_status - pending_updates, last_patch, reboot_pending across full stack
This commit is contained in:
AetherForge
@@ -3,22 +3,278 @@ package api
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"crypto-miner-server/internal/db"
|
||||
"crypto-miner-server/internal/ollama"
|
||||
)
|
||||
|
||||
func TestHandleReportSingleAndArray(t *testing.T) {
|
||||
// Numeric constants from ai_handler.go (guard against silent drift).
|
||||
const (
|
||||
aiReportBufferCap = 1000
|
||||
aiHeartbeatIntervalSeconds = 60
|
||||
aiReasoningTruncateLen = 120
|
||||
aiReportOutputTruncateLen = 200
|
||||
aiEngineIdleTTL = time.Hour
|
||||
aiActivityIdleTTL = 24 * time.Hour
|
||||
aiCleanupTickerInterval = 10 * time.Minute
|
||||
)
|
||||
|
||||
func newTestAIHandler(t *testing.T) *AIHandler {
|
||||
t.Helper()
|
||||
database, err := db.New(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer database.Close()
|
||||
t.Cleanup(func() { database.Close() })
|
||||
return NewAIHandler(database)
|
||||
}
|
||||
|
||||
h := NewAIHandler(database)
|
||||
func mockOllamaChatServer(t *testing.T, content string, status int) *httptest.Server {
|
||||
t.Helper()
|
||||
if status == 0 {
|
||||
status = http.StatusOK
|
||||
}
|
||||
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/api/chat" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
if r.Method != http.MethodPost {
|
||||
http.Error(w, "method", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
w.WriteHeader(status)
|
||||
if status != http.StatusOK {
|
||||
_, _ = w.Write([]byte("ollama down"))
|
||||
return
|
||||
}
|
||||
payload := map[string]interface{}{
|
||||
"message": map[string]string{"content": content},
|
||||
"done": true,
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(payload)
|
||||
}))
|
||||
}
|
||||
|
||||
func ollamaDecideContent(reasoning string, tools []ollama.ToolCall) string {
|
||||
resp := ollama.DecideResponse{Reasoning: reasoning, ToolCalls: tools}
|
||||
b, _ := json.Marshal(resp)
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func TestAIConstants(t *testing.T) {
|
||||
if aiReportBufferCap != 1000 || aiHeartbeatIntervalSeconds != 60 ||
|
||||
aiReasoningTruncateLen != 120 || aiReportOutputTruncateLen != 200 ||
|
||||
aiEngineIdleTTL != time.Hour || aiActivityIdleTTL != 24*time.Hour ||
|
||||
aiCleanupTickerInterval != 10*time.Minute {
|
||||
t.Fatal("ai_handler constants drifted from documented values")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAINewAIHandler(t *testing.T) {
|
||||
h := newTestAIHandler(t)
|
||||
if h.db == nil || h.engines == nil || h.activity == nil {
|
||||
t.Fatal("NewAIHandler did not initialize fields")
|
||||
}
|
||||
if cap(h.reports) != aiReportBufferCap {
|
||||
t.Fatalf("reports slice cap want %d got %d", aiReportBufferCap, cap(h.reports))
|
||||
}
|
||||
if len(h.ActivitySnapshot()) != 0 {
|
||||
t.Fatal("expected empty activity snapshot")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAITruncateStr(t *testing.T) {
|
||||
if truncateStr("short", aiReasoningTruncateLen) != "short" {
|
||||
t.Fatal("short string should be unchanged")
|
||||
}
|
||||
long := strings.Repeat("x", aiReasoningTruncateLen+10)
|
||||
out := truncateStr(long, aiReasoningTruncateLen)
|
||||
if len(out) != aiReasoningTruncateLen+3 || !strings.HasSuffix(out, "...") {
|
||||
t.Fatalf("truncate want len %d with ellipsis, got %q", aiReasoningTruncateLen+3, out)
|
||||
}
|
||||
if truncateStr("", 5) != "" {
|
||||
t.Fatal("empty string")
|
||||
}
|
||||
if truncateStr("exact", 5) != "exact" {
|
||||
t.Fatal("exact length boundary")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAISetEngineForAgentAndGetRemove(t *testing.T) {
|
||||
h := newTestAIHandler(t)
|
||||
srv := mockOllamaChatServer(t, ollamaDecideContent("ok", nil), http.StatusOK)
|
||||
defer srv.Close()
|
||||
|
||||
h.SetEngineForAgent("agent-a", srv.URL, "test-model")
|
||||
if h.GetEngine("agent-a") == nil {
|
||||
t.Fatal("GetEngine returned nil after SetEngineForAgent")
|
||||
}
|
||||
h.RemoveEngine("agent-a")
|
||||
if h.GetEngine("agent-a") != nil {
|
||||
t.Fatal("GetEngine should be nil after RemoveEngine")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAISetEngineForAgentDefaults(t *testing.T) {
|
||||
h := newTestAIHandler(t)
|
||||
var gotModel string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var req struct {
|
||||
Model string `json:"model"`
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&req)
|
||||
gotModel = req.Model
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_ = json.NewEncoder(w).Encode(map[string]interface{}{
|
||||
"message": map[string]string{"content": ollamaDecideContent("ok", nil)},
|
||||
"done": true,
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
h.SetEngineForAgent("agent-defaults", srv.URL, "")
|
||||
eng := h.GetEngine("agent-defaults")
|
||||
if eng == nil {
|
||||
t.Fatal("engine missing")
|
||||
}
|
||||
_, err := eng.Decide(&ollama.AgentState{AgentID: "agent-defaults"})
|
||||
if err != nil {
|
||||
t.Fatalf("Decide: %v", err)
|
||||
}
|
||||
if gotModel != "llama3.2" {
|
||||
t.Fatalf("default model want llama3.2 got %q", gotModel)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIHandleDecideErrors(t *testing.T) {
|
||||
h := newTestAIHandler(t)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/decide", strings.NewReader("{"))
|
||||
w := httptest.NewRecorder()
|
||||
h.HandleDecide(w, req)
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatalf("invalid JSON status %d", w.Code)
|
||||
}
|
||||
|
||||
req = httptest.NewRequest(http.MethodPost, "/api/v1/agent/decide", strings.NewReader(`{"worker_name":"x"}`))
|
||||
w = httptest.NewRecorder()
|
||||
h.HandleDecide(w, req)
|
||||
if w.Code != http.StatusBadRequest || !strings.Contains(w.Body.String(), "agent_id") {
|
||||
t.Fatalf("missing agent_id status %d body %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIHandleDecideSuccess(t *testing.T) {
|
||||
h := newTestAIHandler(t)
|
||||
content := ollamaDecideContent("all good", []ollama.ToolCall{{
|
||||
Tool: "check_miner", Args: map[string]string{"process_name": "xmrig"}, Reason: "verify",
|
||||
}})
|
||||
srv := mockOllamaChatServer(t, content, http.StatusOK)
|
||||
defer srv.Close()
|
||||
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"agent_id": "agent-decide-ok",
|
||||
"ollama_endpoint": srv.URL,
|
||||
"model": "test-model",
|
||||
"hostname": "host1",
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/decide", bytes.NewReader(body))
|
||||
w := httptest.NewRecorder()
|
||||
h.HandleDecide(w, req)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status %d body %s", w.Code, w.Body.String())
|
||||
}
|
||||
var resp ollama.DecideResponse
|
||||
if err := json.NewDecoder(w.Body).Decode(&resp); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(resp.ToolCalls) != 1 || resp.ToolCalls[0].Tool != "check_miner" {
|
||||
t.Fatalf("unexpected tool calls: %+v", resp.ToolCalls)
|
||||
}
|
||||
|
||||
snap := h.ActivitySnapshot()
|
||||
var found bool
|
||||
for _, e := range snap {
|
||||
if e.AgentID == "agent-decide-ok" && e.LastAction == "check_miner" && e.ToolCallCount == 1 {
|
||||
found = true
|
||||
if len(e.LastReasoning) > aiReasoningTruncateLen+3 {
|
||||
t.Fatalf("reasoning should be truncated to %d", aiReasoningTruncateLen)
|
||||
}
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("activity not recorded: %+v", snap)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIHandleDecideOllamaFailureFallback(t *testing.T) {
|
||||
h := newTestAIHandler(t)
|
||||
srv := mockOllamaChatServer(t, "", http.StatusInternalServerError)
|
||||
defer srv.Close()
|
||||
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"agent_id": "agent-decide-fail",
|
||||
"ollama_endpoint": srv.URL,
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/decide", bytes.NewReader(body))
|
||||
w := httptest.NewRecorder()
|
||||
h.HandleDecide(w, req)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("fallback should return 200, got %d", w.Code)
|
||||
}
|
||||
var out map[string]interface{}
|
||||
if err := json.NewDecoder(w.Body).Decode(&out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if out["error"] == nil {
|
||||
t.Fatal("expected error field")
|
||||
}
|
||||
tools, ok := out["tool_calls"].([]interface{})
|
||||
if !ok || len(tools) != 1 {
|
||||
t.Fatalf("expected sleep fallback tool_calls, got %v", out["tool_calls"])
|
||||
}
|
||||
first, _ := tools[0].(map[string]interface{})
|
||||
if first["tool"] != "sleep" {
|
||||
t.Fatalf("expected sleep tool, got %v", first["tool"])
|
||||
}
|
||||
args, _ := first["args"].(map[string]interface{})
|
||||
if args["seconds"] != "120" {
|
||||
t.Fatalf("fallback sleep seconds want 120 got %v", args["seconds"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIHandleDecideCreatesEngineOnFirstRequest(t *testing.T) {
|
||||
h := newTestAIHandler(t)
|
||||
content := ollamaDecideContent("idle", nil)
|
||||
srv := mockOllamaChatServer(t, content, http.StatusOK)
|
||||
defer srv.Close()
|
||||
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"agent_id": "agent-new",
|
||||
"ollama_endpoint": srv.URL,
|
||||
"model": "m1",
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/decide", bytes.NewReader(body))
|
||||
w := httptest.NewRecorder()
|
||||
h.HandleDecide(w, req)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status %d %s", w.Code, w.Body.String())
|
||||
}
|
||||
if h.GetEngine("agent-new") == nil {
|
||||
t.Fatal("engine should exist after first decide")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIHandleReportSingleAndArray(t *testing.T) {
|
||||
h := newTestAIHandler(t)
|
||||
|
||||
single := ollama.Report{
|
||||
AgentID: "agent-1",
|
||||
@@ -54,6 +310,166 @@ func TestHandleReportSingleAndArray(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIHandleReportErrors(t *testing.T) {
|
||||
h := newTestAIHandler(t)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/report", errReader{})
|
||||
w := httptest.NewRecorder()
|
||||
h.HandleReport(w, req)
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatalf("read error status %d", w.Code)
|
||||
}
|
||||
|
||||
req = httptest.NewRequest(http.MethodPost, "/api/v1/agent/report", strings.NewReader("not-json"))
|
||||
w = httptest.NewRecorder()
|
||||
h.HandleReport(w, req)
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatalf("invalid json status %d", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIHandleReportBufferCap(t *testing.T) {
|
||||
h := newTestAIHandler(t)
|
||||
for i := 0; i < aiReportBufferCap+5; i++ {
|
||||
r := ollama.Report{AgentID: "cap-agent", Tool: "sleep", Success: true}
|
||||
b, _ := json.Marshal(r)
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/report", bytes.NewReader(b))
|
||||
w := httptest.NewRecorder()
|
||||
h.HandleReport(w, req)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("report %d status %d", i, w.Code)
|
||||
}
|
||||
}
|
||||
h.mu.RLock()
|
||||
n := len(h.reports)
|
||||
h.mu.RUnlock()
|
||||
if n != aiReportBufferCap {
|
||||
t.Fatalf("reports len want %d got %d", aiReportBufferCap, n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIHandleReportEventBroadcaster(t *testing.T) {
|
||||
h := newTestAIHandler(t)
|
||||
var events []AIActivityEntry
|
||||
h.SetEventBroadcaster(func(e AIActivityEntry) {
|
||||
events = append(events, e)
|
||||
})
|
||||
|
||||
r := ollama.Report{AgentID: "evt-agent", Tool: "upload_log", Success: true, Output: "done"}
|
||||
b, _ := json.Marshal(r)
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/report", bytes.NewReader(b))
|
||||
w := httptest.NewRecorder()
|
||||
h.HandleReport(w, req)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status %d", w.Code)
|
||||
}
|
||||
if len(events) != 1 || events[0].AgentID != "evt-agent" || events[0].LastTool != "upload_log" {
|
||||
t.Fatalf("broadcaster events: %+v", events)
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
_ = json.NewDecoder(bytes.NewReader(w.Body.Bytes())).Decode(&resp)
|
||||
if resp["received"].(float64) != 1 || resp["success"] != true {
|
||||
t.Fatalf("response: %v", resp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIHandleHeartbeat(t *testing.T) {
|
||||
h := newTestAIHandler(t)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/heartbeat", strings.NewReader("{"))
|
||||
w := httptest.NewRecorder()
|
||||
h.HandleHeartbeat(w, req)
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatalf("invalid JSON status %d", w.Code)
|
||||
}
|
||||
|
||||
req = httptest.NewRequest(http.MethodPost, "/api/v1/agent/heartbeat", strings.NewReader(`{"status":"alive"}`))
|
||||
w = httptest.NewRecorder()
|
||||
h.HandleHeartbeat(w, req)
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatalf("missing agent_id status %d", w.Code)
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(heartbeatRequest{AgentID: "hb-agent", Status: "alive", Message: "ok"})
|
||||
req = httptest.NewRequest(http.MethodPost, "/api/v1/agent/heartbeat", bytes.NewReader(body))
|
||||
w = httptest.NewRecorder()
|
||||
h.HandleHeartbeat(w, req)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status %d %s", w.Code, w.Body.String())
|
||||
}
|
||||
var out map[string]interface{}
|
||||
if err := json.NewDecoder(w.Body).Decode(&out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if out["success"] != true {
|
||||
t.Fatalf("success: %v", out["success"])
|
||||
}
|
||||
interval, ok := out["interval"].(float64)
|
||||
if !ok || int(interval) != aiHeartbeatIntervalSeconds {
|
||||
t.Fatalf("interval want %d got %v", aiHeartbeatIntervalSeconds, out["interval"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIRecordActivityMergeAndBroadcast(t *testing.T) {
|
||||
h := newTestAIHandler(t)
|
||||
now := time.Now()
|
||||
h.mu.Lock()
|
||||
h.activity["merge-agent"] = AIActivityEntry{
|
||||
AgentID: "merge-agent",
|
||||
LastReportAt: now,
|
||||
LastSuccess: true,
|
||||
ToolCallCount: 3,
|
||||
}
|
||||
h.mu.Unlock()
|
||||
|
||||
var broadcasted AIActivityEntry
|
||||
h.SetEventBroadcaster(func(e AIActivityEntry) {
|
||||
broadcasted = e
|
||||
})
|
||||
|
||||
h.recordActivity(AIActivityEntry{
|
||||
AgentID: "merge-agent",
|
||||
LastDecideAt: now.Add(time.Minute),
|
||||
LastAction: "sleep",
|
||||
LastTool: "sleep",
|
||||
LastReasoning: "rest",
|
||||
ToolCallCount: 0, // zero must not wipe previous count
|
||||
})
|
||||
|
||||
h.mu.RLock()
|
||||
prev := h.activity["merge-agent"]
|
||||
h.mu.RUnlock()
|
||||
if prev.ToolCallCount != 3 {
|
||||
t.Fatalf("zero ToolCallCount should not clear prior count, got %d", prev.ToolCallCount)
|
||||
}
|
||||
if prev.LastAction != "sleep" || prev.LastReasoning != "rest" {
|
||||
t.Fatalf("merged activity: %+v", prev)
|
||||
}
|
||||
if broadcasted.AgentID != "merge-agent" {
|
||||
t.Fatalf("broadcast: %+v", broadcasted)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIActivitySnapshot(t *testing.T) {
|
||||
h := newTestAIHandler(t)
|
||||
h.mu.Lock()
|
||||
h.activity["a"] = AIActivityEntry{AgentID: "a"}
|
||||
h.activity["b"] = AIActivityEntry{AgentID: "b"}
|
||||
h.mu.Unlock()
|
||||
snap := h.ActivitySnapshot()
|
||||
if len(snap) != 2 {
|
||||
t.Fatalf("want 2 entries got %d", len(snap))
|
||||
}
|
||||
}
|
||||
|
||||
type errReader struct{}
|
||||
|
||||
func (errReader) Read([]byte) (int, error) { return 0, errors.New("read fail") }
|
||||
func (errReader) Close() error { return nil }
|
||||
|
||||
// errReader must satisfy io.ReadCloser for request body
|
||||
var _ io.ReadCloser = errReader{}
|
||||
|
||||
func TestEstimateXMRPerDay(t *testing.T) {
|
||||
out := EstimateXMRPerDay(3_000_000_000)
|
||||
xmr, ok := out["xmr_per_day"].(float64)
|
||||
|
||||
Reference in New Issue
Block a user