748 lines
23 KiB
Go
748 lines
23 KiB
Go
package main
|
|
|
|
// hub_extra_test.go — session_info, spawn_status push, send/overflow
|
|
// branches, upgrade failures, closed-store resilience.
|
|
|
|
import (
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
func TestHubSessionInfoUpdatesSession(t *testing.T) {
|
|
ts, store := newTestServer(t)
|
|
ws := dialAgent(t, ts)
|
|
_ = ws.WriteJSON(helloFrame("s1"))
|
|
_ = readFrame(t, ws)
|
|
|
|
// malformed session_info (no session payload) is ignored, conn stays up
|
|
_ = ws.WriteJSON(map[string]any{"v": 1, "type": evSessionInfo, "sessionId": "s1", "seq": 1, "ts": 2})
|
|
renamed := "renamed"
|
|
// envelope sessionId missing → falls back to session.id
|
|
info := map[string]any{"id": "s1", "name": renamed, "cwd": "/w", "model": "glm-4", "provider": "p", "startedAt": 99}
|
|
_ = ws.WriteJSON(map[string]any{"v": 1, "type": evSessionInfo, "seq": 2, "ts": 3, "session": info})
|
|
|
|
waitFor(t, 2*time.Second, func() bool {
|
|
rows, err := store.Sessions()
|
|
if err != nil || len(rows) != 1 {
|
|
return false
|
|
}
|
|
return rows[0].Info.Model == "glm-4"
|
|
})
|
|
rows, _ := store.Sessions()
|
|
if rows[0].Info.ID != "s1" || rows[0].Info.Name == nil || *rows[0].Info.Name != renamed {
|
|
t.Fatalf("session after session_info = %+v", rows[0].Info)
|
|
}
|
|
}
|
|
|
|
func TestHubHelloFallbacksAndBadPayload(t *testing.T) {
|
|
ts, _ := newTestServer(t)
|
|
ws := dialAgent(t, ts)
|
|
|
|
// hello without session payload → logged, ignored, conn lives on
|
|
_ = ws.WriteJSON(map[string]any{"v": 1, "type": evHello, "sessionId": "s1", "seq": 0, "ts": 1})
|
|
// hello with session.id but no envelope sessionId → registers under session.id
|
|
frame := helloFrame("ignored")
|
|
frame["sessionId"] = nil
|
|
_ = ws.WriteJSON(frame)
|
|
welcome := readFrame(t, ws)
|
|
if welcome["type"] != evWelcome || welcome["sessionId"] != "ignored" {
|
|
t.Fatalf("welcome = %v, want fallback to session.id", welcome)
|
|
}
|
|
}
|
|
|
|
func TestHubIsOnline(t *testing.T) {
|
|
hub := NewHub(openTestStore(t))
|
|
if hub.IsOnline("nope") {
|
|
t.Fatal("unknown session cannot be online")
|
|
}
|
|
ac := &agentConn{hub: hub, conn: throwawayConn(t), send: make(chan []byte, 1), done: make(chan struct{})}
|
|
hub.agents["s1"] = ac
|
|
if !hub.IsOnline("s1") {
|
|
t.Fatal("registered conn must report online")
|
|
}
|
|
}
|
|
|
|
func TestHubBroadcastSpawnStatus(t *testing.T) {
|
|
ts, _, hub := newTestServerHub(t)
|
|
hub.SpawnStatus = func() []SpawnJob {
|
|
return []SpawnJob{{Repo: "group/proj", State: stateRunning, ContainerID: "cid-1"}}
|
|
}
|
|
web := dialWeb(t, ts)
|
|
if first := readFrame(t, web); first["type"] != frameSessionList {
|
|
t.Fatalf("first frame = %v", first)
|
|
}
|
|
hub.BroadcastSpawnStatus()
|
|
_ = web.SetReadDeadline(time.Now().Add(2 * time.Second))
|
|
var got map[string]any
|
|
for got == nil {
|
|
if m := readFrame(t, web); m["type"] == frameSpawnStatus {
|
|
got = m
|
|
}
|
|
}
|
|
jobs := got["jobs"].([]any)
|
|
if len(jobs) != 1 || jobs[0].(map[string]any)["repo"] != "group/proj" {
|
|
t.Fatalf("spawn_status frame = %v", got)
|
|
}
|
|
}
|
|
|
|
// newTestServerHub is newTestServer but also hands back the hub so tests can
|
|
// wire SpawnStatus themselves (NewServer normally does).
|
|
func newTestServerHub(t *testing.T) (*httptest.Server, *Store, *Hub) {
|
|
t.Helper()
|
|
store := openTestStore(t)
|
|
daemonToken = testToken
|
|
hub := NewHub(store)
|
|
srv := &Server{store: store, hub: hub, gitlab: NewGitLab(store, "https://gitlab.example")}
|
|
ts := httptest.NewServer(srv.Routes(""))
|
|
t.Cleanup(ts.Close)
|
|
return ts, store, hub
|
|
}
|
|
|
|
// throwawayConn is a client-side websocket whose only job is being closeable
|
|
// (agentConn.drop touches conn).
|
|
func throwawayConn(t *testing.T) *websocket.Conn {
|
|
t.Helper()
|
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
c, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
|
|
if err != nil {
|
|
return
|
|
}
|
|
_ = c.Close()
|
|
}))
|
|
t.Cleanup(up.Close)
|
|
conn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(up.URL, "http"), nil)
|
|
if err != nil {
|
|
t.Fatalf("dial throwaway: %v", err)
|
|
}
|
|
t.Cleanup(func() { _ = conn.Close() })
|
|
return conn
|
|
}
|
|
|
|
func dialWeb(t *testing.T, ts *httptest.Server) *websocket.Conn {
|
|
t.Helper()
|
|
wsURL := "ws" + strings.TrimPrefix(ts.URL, "http") + "/ws?token=" + testToken
|
|
web, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
|
if err != nil {
|
|
t.Fatalf("dial web ws: %v", err)
|
|
}
|
|
t.Cleanup(func() { _ = web.Close() })
|
|
return web
|
|
}
|
|
|
|
func TestHubUpgradeFailures(t *testing.T) {
|
|
ts, _ := newTestServer(t)
|
|
|
|
get := func(path string) int {
|
|
req, _ := http.NewRequest(http.MethodGet, ts.URL+path, nil)
|
|
req.Header.Set("Authorization", "Bearer "+testToken)
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatalf("get %s: %v", path, err)
|
|
}
|
|
resp.Body.Close()
|
|
return resp.StatusCode
|
|
}
|
|
// plain HTTP GETs (no upgrade headers) must not 500
|
|
if code := get("/agent/ws"); code != http.StatusBadRequest {
|
|
t.Fatalf("agent ws non-upgrade = %d, want 400", code)
|
|
}
|
|
if code := get("/ws?token=" + testToken); code != http.StatusBadRequest {
|
|
t.Fatalf("web ws non-upgrade = %d, want 400", code)
|
|
}
|
|
}
|
|
|
|
func TestSendToAgentBranches(t *testing.T) {
|
|
hub := NewHub(openTestStore(t))
|
|
|
|
// dropped conn: done closed → ErrOffline via done branch (send kept full
|
|
// so the send case cannot win the select race)
|
|
dead := &agentConn{hub: hub, conn: throwawayConn(t), send: make(chan []byte, 1), done: make(chan struct{})}
|
|
dead.send <- []byte("fill")
|
|
dead.drop()
|
|
hub.agents["gone"] = dead
|
|
if err := hub.sendToAgent("gone", map[string]any{"type": evAbort}); err != ErrOffline {
|
|
t.Fatalf("sendToAgent dropped conn = %v, want ErrOffline", err)
|
|
}
|
|
|
|
// live conn but nobody drains send: timeout branch (3s)
|
|
stuck := &agentConn{hub: hub, conn: throwawayConn(t), send: make(chan []byte, 1), done: make(chan struct{})}
|
|
stuck.send <- []byte("fill")
|
|
hub.agents["stuck"] = stuck
|
|
start := time.Now()
|
|
if err := hub.sendToAgent("stuck", map[string]any{"type": evAbort}); err != ErrOffline {
|
|
t.Fatalf("sendToAgent stuck conn = %v, want ErrOffline", err)
|
|
}
|
|
if elapsed := time.Since(start); elapsed < promptSendTimeout-100*time.Millisecond {
|
|
t.Fatalf("sendToAgent returned after %v, want full timeout", elapsed)
|
|
}
|
|
}
|
|
|
|
func TestWebClientOverflowDrops(t *testing.T) {
|
|
// real ws conn so drop() has something to close
|
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
c, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
|
|
if err != nil {
|
|
return
|
|
}
|
|
_ = c.Close()
|
|
}))
|
|
t.Cleanup(up.Close)
|
|
wsURL := "ws" + strings.TrimPrefix(up.URL, "http")
|
|
clientConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
|
if err != nil {
|
|
t.Fatalf("dial: %v", err)
|
|
}
|
|
defer clientConn.Close()
|
|
|
|
hub := NewHub(openTestStore(t))
|
|
c := &webClient{hub: hub, conn: clientConn, done: make(chan struct{})}
|
|
for i := 0; i < webMaxPending; i++ {
|
|
c.deliverEvent(pendingEvent{sessionID: "s", seq: int64(i), raw: []byte("{}")})
|
|
}
|
|
c.deliverEvent(pendingEvent{sessionID: "s", seq: 1, raw: []byte("{}")}) // overflow → dropped
|
|
select {
|
|
case <-c.done:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("overflowed web client was not dropped")
|
|
}
|
|
if !c.dropped {
|
|
t.Fatal("client must be marked dropped")
|
|
}
|
|
// deliveries to a dropped client are no-ops
|
|
c.deliverControl([]byte("{}"))
|
|
if len(c.control) != 0 {
|
|
t.Fatal("dropped client must not queue control frames")
|
|
}
|
|
|
|
// control-frame overflow path on a fresh client
|
|
c2 := &webClient{hub: hub, conn: clientConn, done: make(chan struct{})}
|
|
for i := 0; i < webMaxPending; i++ {
|
|
c2.deliverControl([]byte("{}"))
|
|
}
|
|
c2.deliverControl([]byte("{}"))
|
|
select {
|
|
case <-c2.done:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("control-overflowed web client was not dropped")
|
|
}
|
|
}
|
|
|
|
func TestAgentWritePumpDropsOnWriteError(t *testing.T) {
|
|
serverConn := make(chan *websocket.Conn, 1)
|
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
c, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
|
|
if err != nil {
|
|
return
|
|
}
|
|
serverConn <- c
|
|
for { // keep server side alive until client vanishes
|
|
if _, _, err := c.ReadMessage(); err != nil {
|
|
return
|
|
}
|
|
}
|
|
}))
|
|
t.Cleanup(up.Close)
|
|
clientConn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(up.URL, "http"), nil)
|
|
if err != nil {
|
|
t.Fatalf("dial: %v", err)
|
|
}
|
|
t.Cleanup(func() { _ = clientConn.Close() })
|
|
var server *websocket.Conn
|
|
select {
|
|
case server = <-serverConn:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("server conn never upgraded")
|
|
}
|
|
|
|
hub := NewHub(openTestStore(t))
|
|
ac := &agentConn{hub: hub, conn: server, send: make(chan []byte, 4), done: make(chan struct{})}
|
|
go ac.writePump()
|
|
// healthy write first
|
|
ac.send <- []byte(`{"type":"welcome"}`)
|
|
_ = clientConn.SetReadDeadline(time.Now().Add(2 * time.Second))
|
|
if _, _, err := clientConn.ReadMessage(); err != nil {
|
|
t.Fatalf("first write lost: %v", err)
|
|
}
|
|
// kill the TCP conn underneath gorilla, then keep writing
|
|
if err := clientConn.UnderlyingConn().Close(); err != nil {
|
|
t.Fatalf("close tcp: %v", err)
|
|
}
|
|
deadline := time.After(3 * time.Second)
|
|
for {
|
|
select {
|
|
case ac.send <- []byte(`{"type":"abort"}`):
|
|
case <-ac.done:
|
|
return // writePump hit the write error and dropped
|
|
case <-deadline:
|
|
t.Fatal("writePump did not drop after write failure")
|
|
}
|
|
time.Sleep(10 * time.Millisecond)
|
|
}
|
|
}
|
|
|
|
func TestHubClosedStoreKeepsServing(t *testing.T) {
|
|
ts, store := newTestServer(t)
|
|
if err := store.Close(); err != nil {
|
|
t.Fatalf("close store: %v", err)
|
|
}
|
|
|
|
ws := dialAgent(t, ts)
|
|
_ = ws.WriteJSON(helloFrame("s1"))
|
|
// welcome still arrives with lastSeq 0 despite store failures
|
|
if welcome := readFrame(t, ws); welcome["type"] != evWelcome || welcome["lastSeq"].(float64) != 0 {
|
|
t.Fatalf("welcome with broken store = %v", welcome)
|
|
}
|
|
// events are not persisted but fan-out still happens; conn must survive
|
|
_ = ws.WriteJSON(map[string]any{"v": 1, "type": evAgentSettled, "sessionId": "s1", "seq": 1, "ts": 5})
|
|
|
|
hub := NewHub(store) // store already closed
|
|
if views := hub.SessionsView(); len(views) != 0 {
|
|
t.Fatalf("SessionsView with broken store = %v, want empty", views)
|
|
}
|
|
}
|
|
|
|
func TestHubBroadcastSpawnStatusNilProvider(t *testing.T) {
|
|
hub := NewHub(openTestStore(t))
|
|
hub.BroadcastSpawnStatus() // SpawnStatus nil → no-op, must not panic
|
|
hub.SpawnStatus = func() []SpawnJob { return nil }
|
|
hub.BroadcastSpawnStatus() // provider returning nil jobs → empty frame
|
|
}
|
|
|
|
func TestHubByeClosesConnAndEventWithoutSessionIgnored(t *testing.T) {
|
|
ts, _ := newTestServer(t)
|
|
ws := dialAgent(t, ts)
|
|
_ = ws.WriteJSON(helloFrame("s1"))
|
|
_ = readFrame(t, ws)
|
|
|
|
// event frames without sessionId are dropped silently
|
|
_ = ws.WriteJSON(map[string]any{"v": 1, "type": evMessageEnd, "seq": 9, "ts": 9})
|
|
// bye closes the conn from the server side
|
|
_ = ws.WriteJSON(map[string]any{"v": 1, "type": evBye, "sessionId": "s1", "seq": 10, "ts": 10})
|
|
_ = ws.SetReadDeadline(time.Now().Add(2 * time.Second))
|
|
if _, _, err := ws.ReadMessage(); err == nil {
|
|
t.Fatal("conn must close after bye")
|
|
}
|
|
}
|
|
|
|
func TestHubHelloWelcomeSendDropWhenQueueFull(t *testing.T) {
|
|
hub := NewHub(openTestStore(t))
|
|
ac := &agentConn{hub: hub, conn: throwawayConn(t), send: make(chan []byte, 1), done: make(chan struct{})}
|
|
ac.send <- []byte("full") // nobody drains → welcome select takes default → drop
|
|
hello := helloFrame("s1")
|
|
raw, _ := json.Marshal(hello)
|
|
hub.handleHello(ac, frame{typ: evHello, sessionID: "s1", raw: raw})
|
|
select {
|
|
case <-ac.done:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("agent with full send queue must be dropped at hello")
|
|
}
|
|
}
|
|
|
|
func TestHubWebReadWritePumpEdges(t *testing.T) {
|
|
ts, _, hub := newTestServerHub(t)
|
|
web := dialWeb(t, ts)
|
|
if first := readFrame(t, web); first["type"] != frameSessionList {
|
|
t.Fatalf("first frame = %v", first)
|
|
}
|
|
// garbage frames must not kill the read pump
|
|
if err := web.WriteMessage(websocket.TextMessage, []byte("not json")); err != nil {
|
|
t.Fatalf("write garbage: %v", err)
|
|
}
|
|
if err := web.WriteJSON(map[string]any{"type": frameSubscribe, "sessionId": "s1"}); err != nil {
|
|
t.Fatalf("subscribe: %v", err)
|
|
}
|
|
waitFor(t, 2*time.Second, func() bool {
|
|
hub.mu.Lock()
|
|
defer hub.mu.Unlock()
|
|
for c := range hub.webs {
|
|
return c.subscription() == "s1"
|
|
}
|
|
return false
|
|
})
|
|
if err := web.WriteJSON(map[string]any{"type": frameUnsubscribe}); err != nil {
|
|
t.Fatalf("unsubscribe: %v", err)
|
|
}
|
|
waitFor(t, 2*time.Second, func() bool {
|
|
hub.mu.Lock()
|
|
defer hub.mu.Unlock()
|
|
for c := range hub.webs {
|
|
return c.subscription() == ""
|
|
}
|
|
return false
|
|
})
|
|
// client disconnects → readPump exits, client unregistered
|
|
_ = web.Close()
|
|
waitFor(t, 2*time.Second, func() bool {
|
|
hub.mu.Lock()
|
|
defer hub.mu.Unlock()
|
|
return len(hub.webs) == 0
|
|
})
|
|
}
|
|
|
|
func TestHubDropWebIdempotent(t *testing.T) {
|
|
hub := NewHub(openTestStore(t))
|
|
c := &webClient{hub: hub, conn: throwawayConn(t), done: make(chan struct{})}
|
|
hub.mu.Lock()
|
|
hub.webs[c] = struct{}{}
|
|
hub.mu.Unlock()
|
|
hub.dropWeb(c)
|
|
hub.dropWeb(c) // second drop hits the already-removed early return
|
|
select {
|
|
case <-c.done:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("web client was not dropped")
|
|
}
|
|
hub.mu.Lock()
|
|
left := len(hub.webs)
|
|
hub.mu.Unlock()
|
|
if left != 0 {
|
|
t.Fatalf("webs after drop = %d, want 0", left)
|
|
}
|
|
}
|
|
|
|
func TestWebWritePumpDropsOnWriteError(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name string
|
|
ctrl bool // deliver control frames (vs events only)
|
|
}{
|
|
{"control-write", true},
|
|
{"events-write", false},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
serverConn := make(chan *websocket.Conn, 1)
|
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
c, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
|
|
if err != nil {
|
|
return
|
|
}
|
|
serverConn <- c
|
|
for {
|
|
if _, _, err := c.ReadMessage(); err != nil {
|
|
return
|
|
}
|
|
}
|
|
}))
|
|
t.Cleanup(up.Close)
|
|
clientConn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(up.URL, "http"), nil)
|
|
if err != nil {
|
|
t.Fatalf("dial: %v", err)
|
|
}
|
|
t.Cleanup(func() { _ = clientConn.Close() })
|
|
var server *websocket.Conn
|
|
select {
|
|
case server = <-serverConn:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("server conn never upgraded")
|
|
}
|
|
|
|
hub := NewHub(openTestStore(t))
|
|
c := &webClient{hub: hub, conn: server, done: make(chan struct{})}
|
|
hub.mu.Lock()
|
|
hub.webs[c] = struct{}{} // writePump errors route through dropWeb
|
|
hub.mu.Unlock()
|
|
go c.writePump()
|
|
if tc.ctrl {
|
|
c.deliverControl([]byte(`{"type":"session_list"}`))
|
|
} else {
|
|
c.deliverEvent(pendingEvent{sessionID: "s", seq: 1, raw: []byte("{}")})
|
|
}
|
|
_ = clientConn.SetReadDeadline(time.Now().Add(2 * time.Second))
|
|
if _, _, err := clientConn.ReadMessage(); err != nil {
|
|
t.Fatalf("first frame lost: %v", err)
|
|
}
|
|
if err := clientConn.UnderlyingConn().Close(); err != nil {
|
|
t.Fatalf("close tcp: %v", err)
|
|
}
|
|
// flush cycle hits the dead conn → writePump drops the client
|
|
deadline := time.After(3 * time.Second)
|
|
for {
|
|
if tc.ctrl {
|
|
c.deliverControl([]byte(`{"type":"session_list"}`))
|
|
} else {
|
|
c.deliverEvent(pendingEvent{sessionID: "s", seq: 2, raw: []byte("{}")})
|
|
}
|
|
select {
|
|
case <-c.done:
|
|
return
|
|
case <-deadline:
|
|
t.Fatal("writePump did not drop after write failure")
|
|
default:
|
|
}
|
|
time.Sleep(2 * time.Millisecond)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestWebClientDeliverToDropped(t *testing.T) {
|
|
hub := NewHub(openTestStore(t))
|
|
c := &webClient{hub: hub, conn: throwawayConn(t), done: make(chan struct{})}
|
|
for i := 0; i < webMaxPending+1; i++ { // overflow sets the dropped flag
|
|
c.deliverEvent(pendingEvent{sessionID: "s", seq: int64(i), raw: []byte("{}")})
|
|
}
|
|
c.deliverEvent(pendingEvent{sessionID: "s", seq: 1, raw: []byte("{}")})
|
|
if len(c.events) != webMaxPending {
|
|
t.Fatalf("dropped client queued %d events, want %d", len(c.events), webMaxPending)
|
|
}
|
|
}
|
|
|
|
func TestWebClientByteOverflowDrops(t *testing.T) {
|
|
hub := NewHub(openTestStore(t))
|
|
c := &webClient{hub: hub, conn: throwawayConn(t), done: make(chan struct{})}
|
|
// stay under the frame-count cap; cross the byte cap on the last event.
|
|
meg := make([]byte, 1<<20)
|
|
var i int
|
|
for ; (i+1)*(1<<20) <= webMaxPendingByte; i++ {
|
|
c.deliverEvent(pendingEvent{sessionID: "s", seq: int64(i), raw: meg})
|
|
}
|
|
if c.dropped {
|
|
t.Fatalf("client dropped at %dMiB, want only beyond %dMiB", i, webMaxPendingByte>>20)
|
|
}
|
|
c.deliverEvent(pendingEvent{sessionID: "s", seq: int64(i), raw: meg}) // one over the cap
|
|
select {
|
|
case <-c.done:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("byte-overflowed web client was not dropped")
|
|
}
|
|
if !c.dropped {
|
|
t.Fatal("client must be marked dropped on byte overflow")
|
|
}
|
|
}
|
|
|
|
func TestUnregisterReplacedConnKeepsOnline(t *testing.T) {
|
|
store := openTestStore(t)
|
|
hub := NewHub(store)
|
|
if err := store.UpsertSession(SessionInfo{ID: "s1", Cwd: "/w", Model: "m", Provider: "p"}); err != nil {
|
|
t.Fatalf("upsert: %v", err)
|
|
}
|
|
if err := store.SetOnline("s1", true); err != nil {
|
|
t.Fatalf("online: %v", err)
|
|
}
|
|
old := &agentConn{hub: hub, conn: throwawayConn(t), send: make(chan []byte, 1), done: make(chan struct{})}
|
|
old.sessionID = "s1"
|
|
fresh := &agentConn{hub: hub, conn: throwawayConn(t), send: make(chan []byte, 1), done: make(chan struct{})}
|
|
fresh.sessionID = "s1"
|
|
hub.agents["s1"] = old
|
|
hub.agents["s1"] = fresh // reconnect replaced old
|
|
|
|
hub.unregister(old)
|
|
rows, err := store.Sessions()
|
|
if err != nil || len(rows) != 1 {
|
|
t.Fatalf("sessions: %v %d", err, len(rows))
|
|
}
|
|
if !rows[0].OnlineDB {
|
|
t.Fatal("replaced conn unregistering must not flip the live session offline")
|
|
}
|
|
|
|
hub.unregister(fresh)
|
|
rows, err = store.Sessions()
|
|
if err != nil {
|
|
t.Fatalf("sessions: %v", err)
|
|
}
|
|
if rows[0].OnlineDB {
|
|
t.Fatal("last live conn unregistering must mark the session offline")
|
|
}
|
|
}
|
|
|
|
func TestPingPongConstants(t *testing.T) {
|
|
if pingPeriod >= pongWait {
|
|
t.Fatalf("pingPeriod %v must be < pongWait %v", pingPeriod, pongWait)
|
|
}
|
|
}
|
|
|
|
func TestWebClientSendPing(t *testing.T) {
|
|
serverConn := make(chan *websocket.Conn, 1)
|
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
c, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
|
|
if err != nil {
|
|
return
|
|
}
|
|
serverConn <- c
|
|
for {
|
|
if _, _, err := c.ReadMessage(); err != nil {
|
|
return
|
|
}
|
|
}
|
|
}))
|
|
t.Cleanup(up.Close)
|
|
clientConn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(up.URL, "http"), nil)
|
|
if err != nil {
|
|
t.Fatalf("dial: %v", err)
|
|
}
|
|
t.Cleanup(func() { _ = clientConn.Close() })
|
|
var server *websocket.Conn
|
|
select {
|
|
case server = <-serverConn:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("server conn never upgraded")
|
|
}
|
|
hub := NewHub(openTestStore(t))
|
|
c := &webClient{hub: hub, conn: server, done: make(chan struct{})}
|
|
if !c.sendPing() {
|
|
t.Fatal("sendPing over a live conn must succeed")
|
|
}
|
|
ac := &agentConn{hub: hub, conn: server, send: make(chan []byte, 1), done: make(chan struct{})}
|
|
if !ac.sendPing() {
|
|
t.Fatal("agent sendPing over a live conn must succeed")
|
|
}
|
|
}
|
|
|
|
func TestPongFramesKeepReadLoopsAlive(t *testing.T) {
|
|
// unsolicited pongs must pass through both read pumps (refreshing the
|
|
// read deadline) without killing the conn.
|
|
ts, _ := newTestServer(t)
|
|
ws := dialAgent(t, ts)
|
|
if err := ws.WriteMessage(websocket.PongMessage, nil); err != nil {
|
|
t.Fatalf("agent pong: %v", err)
|
|
}
|
|
_ = ws.WriteJSON(helloFrame("s1"))
|
|
if welcome := readFrame(t, ws); welcome["type"] != evWelcome {
|
|
t.Fatalf("agent read loop died after pong: %v", welcome)
|
|
}
|
|
|
|
web := dialWeb(t, ts)
|
|
if first := readFrame(t, web); first["type"] != frameSessionList {
|
|
t.Fatalf("first web frame = %v", first)
|
|
}
|
|
if err := web.WriteMessage(websocket.PongMessage, nil); err != nil {
|
|
t.Fatalf("web pong: %v", err)
|
|
}
|
|
// subscribe processed by the (still live) read pump → live events flow.
|
|
if err := web.WriteJSON(map[string]any{"type": frameSubscribe, "sessionId": "s1"}); err != nil {
|
|
t.Fatalf("subscribe: %v", err)
|
|
}
|
|
// A reader goroutine avoids deadline-based reads (gorilla conns fail
|
|
// permanently after a read timeout); resend until a batch arrives.
|
|
frames := make(chan map[string]any, 16)
|
|
go func() {
|
|
defer close(frames)
|
|
for {
|
|
var m map[string]any
|
|
if err := web.ReadJSON(&m); err != nil {
|
|
return
|
|
}
|
|
frames <- m
|
|
}
|
|
}()
|
|
start := time.Now()
|
|
for i := 0; ; i++ {
|
|
_ = ws.WriteJSON(map[string]any{"v": 1, "type": evMessageUpdate, "sessionId": "s1", "seq": int64(i + 1), "ts": 10, "delta": "x"})
|
|
select {
|
|
case m, ok := <-frames:
|
|
if !ok {
|
|
t.Fatal("web conn closed; read loop died after pong")
|
|
}
|
|
if m["type"] == frameEvents {
|
|
return // read pump survived the pong and routed the subscription
|
|
}
|
|
case <-time.After(200 * time.Millisecond):
|
|
}
|
|
if time.Since(start) > 3*time.Second {
|
|
t.Fatal("no events frame delivered after pong")
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestWebResubscribeSplitsEventBatches(t *testing.T) {
|
|
ts, _, hub := newTestServerHub(t)
|
|
agent1 := dialAgent(t, ts)
|
|
_ = agent1.WriteJSON(helloFrame("s1"))
|
|
_ = readFrame(t, agent1)
|
|
agent2 := dialAgent(t, ts)
|
|
_ = agent2.WriteJSON(helloFrame("s2"))
|
|
_ = readFrame(t, agent2)
|
|
|
|
web := dialWeb(t, ts)
|
|
if first := readFrame(t, web); first["type"] != frameSessionList {
|
|
t.Fatalf("first web frame = %v", first)
|
|
}
|
|
|
|
webSubscribed := func(want string) func() bool {
|
|
return func() bool {
|
|
hub.mu.Lock()
|
|
defer hub.mu.Unlock()
|
|
for c := range hub.webs {
|
|
c.mu.Lock()
|
|
sub := c.sub
|
|
c.mu.Unlock()
|
|
if sub == want {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
}
|
|
_ = web.WriteJSON(map[string]any{"type": frameSubscribe, "sessionId": "s1"})
|
|
waitFor(t, 2*time.Second, webSubscribed("s1"))
|
|
_ = agent1.WriteJSON(map[string]any{"v": 1, "type": evMessageUpdate, "sessionId": "s1", "seq": 1, "ts": 10, "delta": "one"})
|
|
|
|
readBatch := func(want string) map[string]any {
|
|
deadline := time.Now().Add(3 * time.Second)
|
|
for time.Now().Before(deadline) {
|
|
_ = web.SetReadDeadline(time.Now().Add(2 * time.Second))
|
|
var m map[string]any
|
|
if err := web.ReadJSON(&m); err != nil {
|
|
t.Fatalf("read frame (want batch %s): %v", want, err)
|
|
}
|
|
if m["type"] != frameEvents {
|
|
continue
|
|
}
|
|
if m["sessionId"] != want {
|
|
t.Fatalf("events frame sessionId = %v, want %q (sessions must not mix)", m["sessionId"], want)
|
|
}
|
|
for _, e := range m["events"].([]any) {
|
|
if e.(map[string]any)["sessionId"] != want {
|
|
t.Fatalf("batch for %s contains event of %v", want, e.(map[string]any)["sessionId"])
|
|
}
|
|
}
|
|
return m
|
|
}
|
|
t.Fatalf("no events frame for %s before timeout", want)
|
|
return nil
|
|
}
|
|
|
|
readBatch("s1") // event A flushed while subscribed to A
|
|
|
|
_ = web.WriteJSON(map[string]any{"type": frameSubscribe, "sessionId": "s2"})
|
|
waitFor(t, 2*time.Second, webSubscribed("s2"))
|
|
_ = agent2.WriteJSON(map[string]any{"v": 1, "type": evMessageUpdate, "sessionId": "s2", "seq": 1, "ts": 11, "delta": "two"})
|
|
readBatch("s2") // event B must arrive as its own frame, never mixed with A's
|
|
}
|
|
|
|
func TestHubBusyFlagLifecycle(t *testing.T) {
|
|
ts, _, hub := newTestServerHub(t)
|
|
ws := dialAgent(t, ts)
|
|
defer ws.Close()
|
|
_ = ws.WriteJSON(helloFrame("busy-1"))
|
|
_ = readFrame(t, ws)
|
|
sendEv := func(typ string, seq int64) {
|
|
_ = ws.WriteJSON(map[string]any{
|
|
"v": 1, "sessionId": "busy-1", "seq": seq, "ts": 1, "type": typ,
|
|
})
|
|
}
|
|
sendEv("agent_start", 1)
|
|
waitFor(t, 3*time.Second, func() bool {
|
|
for _, s := range hub.SessionsView() {
|
|
if s.ID == "busy-1" {
|
|
return s.Busy
|
|
}
|
|
}
|
|
return false
|
|
})
|
|
sendEv("agent_settled", 2)
|
|
waitFor(t, 3*time.Second, func() bool {
|
|
for _, s := range hub.SessionsView() {
|
|
if s.ID == "busy-1" {
|
|
return !s.Busy
|
|
}
|
|
}
|
|
return false
|
|
})
|
|
}
|