wslog_test.go

Documentation: github.com/cybertec-postgresql/pgwatch/v5/internal/webserver

     1  package webserver
     2  
     3  import (
     4  	"net/http"
     5  	"net/http/httptest"
     6  	"net/url"
     7  	"strings"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/cybertec-postgresql/pgwatch/v5/internal/log"
    12  	"github.com/gorilla/websocket"
    13  	"github.com/stretchr/testify/assert"
    14  	"github.com/stretchr/testify/require"
    15  )
    16  
    17  func TestServeWsLog_UpgradeError(t *testing.T) {
    18  	ts := &WebUIServer{Logger: log.NewNoopLogger()}
    19  	r := httptest.NewRequest(http.MethodGet, "/wslog", nil)
    20  	w := httptest.NewRecorder()
    21  	// No websocket headers, should fail to upgrade
    22  	ts.serveWsLog(w, r)
    23  	resp := w.Result()
    24  	assert.Equal(t, http.StatusBadRequest, resp.StatusCode)
    25  }
    26  
    27  func TestServeWsLog_Success(t *testing.T) {
    28  	ts := &WebUIServer{Logger: log.Init(log.CmdOpts{LogLevel: "debug"})}
    29  	h := http.HandlerFunc(ts.serveWsLog)
    30  	tsServer := httptest.NewServer(h)
    31  	defer tsServer.Close()
    32  
    33  	u := "ws" + strings.TrimPrefix(tsServer.URL, "http")
    34  	uParsed, _ := url.Parse(u)
    35  	uParsed.Path = "/wslog"
    36  
    37  	ws, resp, err := websocket.DefaultDialer.Dial(uParsed.String(), nil)
    38  	require.NotNil(t, ws)
    39  	require.NoError(t, err)
    40  	assert.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode)
    41  
    42  	// send ping message to keep connection alive
    43  	assert.NoError(t, ws.WriteMessage(websocket.PingMessage, nil))
    44  
    45  	// send some log message
    46  	time.Sleep(100 * time.Millisecond)
    47  	ts.Info("Test message")
    48  	// check output though the websocket
    49  	assert.NoError(t, ws.SetReadDeadline(time.Now().Add(2*time.Second)))
    50  	msgType, msg, err := ws.ReadMessage()
    51  	assert.NoError(t, err)
    52  	assert.Equal(t, websocket.TextMessage, msgType)
    53  	assert.NotEmpty(t, msg)
    54  	assert.Contains(t, string(msg), "Test message")
    55  	time.Sleep(4 * time.Second)
    56  	// check if the connection is closed
    57  	assert.NoError(t, ws.Close())
    58  	ts.Info("Test message after websocket close")
    59  	assert.Error(t, ws.SetReadDeadline(time.Now().Add(2*time.Second)))
    60  	_, _, err = ws.ReadMessage()
    61  	assert.Error(t, err, "should error because connection is closed")
    62  }
    63