saerro-go/util/testutil/ws.go
2024-10-28 13:46:52 -07:00

50 lines
822 B
Go

package testutil
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/coder/websocket"
)
type MockESS struct {
Server *httptest.Server
LastMessage string
}
func (m MockESS) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
c, err := websocket.Accept(rw, req, nil)
if err != nil {
rw.WriteHeader(400)
rw.Write([]byte("websocket connection failed"))
return
}
defer c.CloseNow()
ctx, cancel := context.WithTimeout(req.Context(), time.Second*30)
defer cancel()
_, body, err := c.Read(ctx)
if err != nil {
rw.WriteHeader(500)
rw.Write([]byte("websocket read failed"))
return
}
m.LastMessage = string(body)
}
func GetMockESS(t *testing.T) MockESS {
t.Helper()
m := MockESS{}
s := httptest.NewServer(m)
m.Server = s
t.Cleanup(s.Close)
return m
}