Files
AetherForge/server/internal/api/websocket_test.go

979 lines
30 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)
}
var resp Message
if err := conn.ReadJSON(&resp); err != nil {
t.Fatalf("read auth_response: %v", err)
}
return resp
}
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)
}
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)
}
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")
}
}
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",
})
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})
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})
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})
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
}
}