mirror of
https://github.com/rcourtman/Pulse.git
synced 2026-05-06 16:16:26 +00:00
36 lines
1.1 KiB
Go
36 lines
1.1 KiB
Go
package agentexec
|
|
|
|
import "github.com/gorilla/websocket"
|
|
|
|
// TestRegisterAgent registers a fake agent for testing purposes.
|
|
// This bypasses the WebSocket handshake and directly inserts into the agents map.
|
|
// It also installs a no-op writeTextMessage so that send operations don't panic
|
|
// on the nil websocket connection. The original writer is restored via the
|
|
// returned cleanup function (typically deferred by the caller).
|
|
//
|
|
// NOTE: This mutates the package-global writeTextMessage without synchronization.
|
|
// Callers must NOT use t.Parallel() when using this helper.
|
|
func (s *Server) TestRegisterAgent(agentID, hostname string) (cleanup func()) {
|
|
// Save and override the global writer.
|
|
origWriter := writeTextMessage
|
|
writeTextMessage = func(_ *websocket.Conn, _ []byte) error {
|
|
return nil
|
|
}
|
|
|
|
s.mu.Lock()
|
|
s.agents[agentID] = &agentConn{
|
|
agent: ConnectedAgent{
|
|
AgentID: agentID,
|
|
Hostname: hostname,
|
|
},
|
|
done: make(chan struct{}),
|
|
}
|
|
s.mu.Unlock()
|
|
|
|
return func() {
|
|
writeTextMessage = origWriter
|
|
s.mu.Lock()
|
|
delete(s.agents, agentID)
|
|
s.mu.Unlock()
|
|
}
|
|
}
|