mediamtx/internal/protocols/websocket/serverconn_test.go

package websocket

import (
	"context"
	"net"
	"net/http"
	"testing"
	"time"

	"github.com/gorilla/websocket"
	"github.com/stretchr/testify/require"
)

func TestServerConn(t *testing.T) {
	pingReceived := make(chan struct{})
	pingInterval = 100 * time.Millisecond

	s := &http.Server{
		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			c, err := NewServerConn(w, r)
			require.NoError(t, err)
			defer c.Close()

			err = c.WriteJSON("testing")
			require.NoError(t, err)

			<-pingReceived
		}),
	}

	ln, err := net.Listen("tcp", "localhost:6344")
	require.NoError(t, err)

	go s.Serve(ln)
	defer s.Shutdown(context.Background())

	c, res, err := websocket.DefaultDialer.Dial("ws://localhost:6344/", nil)
	require.NoError(t, err)
	defer res.Body.Close()
	defer c.Close() //nolint:errcheck

	c.SetPingHandler(func(_ string) error {
		close(pingReceived)
		return nil
	})

	var msg string
	err = c.ReadJSON(&msg)
	require.NoError(t, err)
	require.Equal(t, "testing", msg)

	_, _, err = c.ReadMessage()
	require.Error(t, err)

	<-pingReceived
}