89 lines
2.3 KiB
Go
89 lines
2.3 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
func TestMaintainConnectionsRestoresClosedClient(t *testing.T) {
|
|
var connections atomic.Int64
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
|
|
if err != nil {
|
|
return
|
|
}
|
|
connections.Add(1)
|
|
defer conn.Close()
|
|
for {
|
|
var packet map[string]json.RawMessage
|
|
if err = conn.ReadJSON(&packet); err != nil {
|
|
return
|
|
}
|
|
for _, raw := range packet {
|
|
var body struct {
|
|
ID string `json:"id"`
|
|
}
|
|
if json.Unmarshal(raw, &body) != nil || body.ID == "" {
|
|
continue
|
|
}
|
|
response := map[string]any{"ctrl": map[string]any{
|
|
"id": body.ID, "code": 200, "text": "ok",
|
|
"params": map[string]any{"user": "usrReconnectTest"},
|
|
}}
|
|
if err = conn.WriteJSON(response); err != nil {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
|
|
account := AccountRecord{
|
|
Machine: "a", Index: 1, Username: "lta0001", Token: "token",
|
|
P2PTopic: "p2pReconnectTest", Endpoint: "ws" + strings.TrimPrefix(server.URL, "http"),
|
|
}
|
|
cfg := &Config{Tenant: "allrich", CommandTimeout: "1s", WSPingInterval: "1h"}
|
|
run := &agentRun{
|
|
cfg: cfg, machine: "a", clients: make(map[int]*WSClient),
|
|
recorder: newDeliveryRecorder("pending", "a", 1), ack: NewDurationHistogram(),
|
|
}
|
|
initial, err := run.connectAccountWithRetry(context.Background(), account, "grpReconnectTest", "api-key", "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
run.setClient(account.Index, initial)
|
|
initial.Close()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
done := make(chan struct{})
|
|
go func() {
|
|
run.maintainConnections(ctx, []AccountRecord{account}, "grpReconnectTest", "api-key", "")
|
|
close(done)
|
|
}()
|
|
deadline := time.Now().Add(3 * time.Second)
|
|
for time.Now().Before(deadline) {
|
|
if replacement := run.client(account.Index); replacement != initial && replacement.alive() {
|
|
cancel()
|
|
<-done
|
|
run.close()
|
|
if got := connections.Load(); got < 2 {
|
|
t.Fatalf("websocket connections = %d, want at least 2", got)
|
|
}
|
|
return
|
|
}
|
|
time.Sleep(10 * time.Millisecond)
|
|
}
|
|
cancel()
|
|
<-done
|
|
run.close()
|
|
t.Fatal("closed client was not restored before deadline")
|
|
}
|