90.7% unit test coverage for ws utility

connection.go

@@ -65,7 +65,7 @@ func getPrefix() string {
 	return ""
 }
 
-func connect(url string, rlConf *readline.Config) []error {
+func connect(url string, rl *readline.Instance) []error {
 	headers := make(http.Header)
 	headers.Add("Origin", options.origin)
 	if options.authHeader != "" {
@@ -86,14 +86,10 @@ func connect(url string, rlConf *readline.Config) []error {
 		return []error{err}
 	}
 	defer func() {
+		rl.Close()
 		TryCloseNormally(ws, "client disconnection")
 		ws.Close()
 	}()
-	rl, err := readline.NewEx(rlConf)
-	if err != nil {
-		return []error{err}
-	}
-	defer rl.Close()
 	session = &Session{
 		ws:     ws,
 		rl:     rl,

connection_test.go

@@ -0,0 +1,154 @@
+package main
+
+import (
+	"context"
+	"io"
+	"os"
+	"regexp"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/chzyer/readline"
+	"github.com/gorilla/websocket"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGetPrefix(t *testing.T) {
+	options.timestamp = true
+	defer func() { options.timestamp = false }()
+	prefix := getPrefix()
+	require.Contains(t, prefix, " ")
+	require.Len(t, prefix, 20)
+	options.timestamp = false
+	prefix = getPrefix()
+	require.Empty(t, prefix)
+}
+
+func TestSession(t *testing.T) {
+	srv := newMockServer(0)
+	defer srv.Close()
+	conn := newMockConn()
+	defer TryCloseNormally(conn, "test finished")
+	ctx, cancel := context.WithCancel(context.Background())
+	outR, outW, _ := os.Pipe()
+	rl, err := readline.NewEx(&readline.Config{Prompt: "> ", Stdout: outW})
+	require.NoError(t, err)
+	s := Session{
+		ws:      conn,
+		rl:      rl,
+		cancel:  cancel,
+		ctx:     ctx,
+		errors:  []error{},
+		errLock: sync.Mutex{},
+	}
+	sent := "test message"
+	// test sendMsg
+	options.timestamp = true
+	defer func() { options.timestamp = false }()
+	err = s.sendMsg(sent)
+	require.NoError(t, err)
+	require.Eventually(t, func() bool { return len(srv.Received) > 0 }, 30*time.Millisecond, 3*time.Millisecond)
+	require.Equal(t, sent, <-srv.Received)
+	// test typing
+	typed := "typed"
+	_, err = rl.WriteStdin([]byte(typed + "\n"))
+	require.NoError(t, err)
+	go func() {
+		s.readConsole()
+	}()
+	require.Eventually(t, func() bool { return len(srv.Received) > 0 }, 20*time.Millisecond, 2*time.Millisecond)
+	require.Equal(t, typed, <-srv.Received)
+	// test readWebsocket
+	go func() {
+		s.readWebsocket()
+	}()
+	// text message
+	srv.ToSend <- sent
+	require.Eventually(t, func() bool { return len(srv.ToSend) == 0 }, 20*time.Millisecond, 2*time.Millisecond)
+	// binary message
+	atomic.StoreInt64(&srv.Mode, websocket.BinaryMessage)
+	srv.ToSend <- sent
+	require.Eventually(t, func() bool { return len(srv.ToSend) == 0 }, 20*time.Millisecond, 2*time.Millisecond)
+	// binary as text
+	options.binAsText = true
+	defer func() { options.binAsText = false }()
+	srv.ToSend <- "binary"
+	require.Eventually(t, func() bool { return len(srv.ToSend) == 0 }, 20*time.Millisecond, 2*time.Millisecond)
+	// filtered
+	toBeFiltered := "must be filtered"
+	options.filter = regexp.MustCompile("^.*not filtered.*$")
+	defer func() { options.filter = nil }()
+	require.False(t, options.filter.MatchString(toBeFiltered))
+	srv.ToSend <- toBeFiltered
+	require.Eventually(t, func() bool { return len(srv.ToSend) == 0 }, 20*time.Millisecond, 2*time.Millisecond)
+	// unknown mode
+	atomic.StoreInt64(&srv.Mode, 0)
+	srv.ToSend <- "unknown"
+	require.Eventually(t, func() bool { return len(srv.ToSend) == 0 }, 20*time.Millisecond, 2*time.Millisecond)
+	time.Sleep(20 * time.Millisecond)
+	cancel()
+	outW.Close()
+	output, err := io.ReadAll(outR)
+	out := string(output)
+	require.NoError(t, err)
+	require.Contains(t, out, " > test message")
+	require.Contains(t, out, " > typed")
+	require.Contains(t, out, " < test message")
+	require.Contains(t, out, " < \n00000000  74 65 73 74 20 6d 65 73  73 61 67 65              |test message|")
+	require.Contains(t, out, " < binary")
+	require.NotContains(t, out, toBeFiltered)
+	require.NotContains(t, out, "unknown")
+	// t.Log(out)
+}
+
+func TestPingPong(t *testing.T) {
+	srv := newMockServer(2 * time.Millisecond)
+	defer srv.Close()
+	options.pingPong = true
+	options.pingInterval = 2 * time.Millisecond
+	outR, outW, _ := os.Pipe()
+	errs := make(chan []error, 1)
+	rl, err := readline.NewEx(&readline.Config{Prompt: "> ", Stdout: outW, UniqueEditLine: true})
+	require.NoError(t, err)
+	// the only way I found to keep redline working for a while
+	go func() {
+		for i := 0; i < 400; i++ {
+			_, err = rl.WriteStdin([]byte("typed"))
+		}
+	}()
+	rl.Write([]byte("typed"))
+	require.NoError(t, err)
+	go func() {
+		errs <- connect(mockURL, rl)
+	}()
+	time.Sleep(200 * time.Millisecond)
+	session.cancel()
+	require.Eventually(t, func() bool { return len(errs) > 0 }, 20*time.Millisecond, 2*time.Millisecond)
+	outW.Close()
+	output, err := io.ReadAll(outR)
+	out := string(output)
+	require.NoError(t, err)
+	require.Contains(t, out, "> ping")
+	require.Contains(t, out, "< pong")
+	require.Contains(t, out, "< ping:")
+	require.Contains(t, out, "> pong:")
+}
+
+func TestInitMsg(t *testing.T) {
+	s := newMockServer(0)
+	defer s.Close()
+	message := "test message"
+	options.initMsg = message
+	defer func() {
+		options.initMsg = ""
+	}()
+	rl, err := readline.New(" >")
+	require.NoError(t, err)
+	time.AfterFunc(500*time.Millisecond, func() { session.cancel() })
+	errs := connect(mockURL, rl)
+	require.Empty(t, errs)
+	require.Eventually(t, func() bool { return len(s.Received) > 0 }, 20*time.Millisecond, 2*time.Millisecond)
+	require.Equal(t, message, <-s.Received)
+}

main.go

@@ -13,7 +13,6 @@ import (
 	"github.com/spf13/cobra"
 )
 
-// Version is app version
 var (
 	version = "local build"
 	options struct {
@@ -89,10 +88,15 @@ func root(cmd *cobra.Command, args []string) {
 	if err == nil {
 		historyFile = filepath.Join(user.HomeDir, ".ws_history")
 	}
-	errs := connect(dest.String(), &readline.Config{
+	rl, err := readline.NewEx(&readline.Config{
 		Prompt:      "> ",
 		HistoryFile: historyFile,
 	})
+	if err != nil {
+		fmt.Fprintln(os.Stderr, err)
+		os.Exit(1)
+	}
+	errs := connect(dest.String(), rl)
 	if len(errs) > 0 {
 		fmt.Println()
 		for _, err := range errs {

main_test.go

@@ -1,95 +1,30 @@
 package main
 
 import (
-	"context"
 	"fmt"
 	"io"
-	"net/url"
 	"os"
 	"os/exec"
 	"strings"
 	"testing"
 	"time"
 
 	"github.com/gorilla/websocket"
-	"github.com/slytomcat/ws/server"
 	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
-const mockURL = "ws://localhost:8080"
-
-type mockServer struct {
-	Close    func() error
-	Received chan string
-	ToSend   chan string
-}
-
-func newMockServer() *mockServer {
-	received := make(chan string, 10)
-	toSend := make(chan string, 10)
-	ctx, cancel := context.WithCancel(context.Background())
-	u, _ := url.Parse(mockURL)
-	s := server.NewServer(u.Host)
-	s.WSHandleFunc("/", func(conn *websocket.Conn) {
-		go func() {
-			for {
-				_, msg, err := conn.ReadMessage()
-				if err != nil {
-					return
-				}
-				received <- string(msg)
-			}
-		}()
-		select {
-		case <-ctx.Done():
-			TryCloseNormally(conn, "server going down")
-			return
-		case data := <-toSend:
-			conn.WriteMessage(websocket.TextMessage, []byte(data))
-		}
-	})
-	go s.ListenAndServe()
-	time.Sleep(50 * time.Millisecond)
-	return &mockServer{
-		Close: func() error {
-			cancel()
-			s.Shutdown(ctx)
-			return nil
-		},
-		Received: received,
-		ToSend:   toSend,
-	}
-}
-
-func newMockConn() *websocket.Conn {
-	dial := websocket.Dialer{}
-	conn, _, err := dial.Dial(mockURL, nil)
-	if err != nil {
-		panic(err)
-	}
-	return conn
-}
-
 func TestMockServer(t *testing.T) {
-	s := newMockServer()
+	s := newMockServer(0)
 	defer s.Close()
 	conn := newMockConn()
 	defer TryCloseNormally(conn, "test finished")
 	sent := "test"
 	// test client -> server message
 	require.NoError(t, conn.WriteMessage(websocket.TextMessage, []byte(sent)))
-	var received string
-	require.Eventually(t, func() bool {
-		select {
-		case received = <-s.Received:
-			return true
-		default:
-			return false
-		}
-	}, 20*time.Millisecond, 2*time.Millisecond)
-	require.Equal(t, sent, received)
+	require.Eventually(t, func() bool { return len(s.Received) > 0 }, 20*time.Millisecond, 2*time.Millisecond)
+	require.Equal(t, sent, <-s.Received)
 	// test server -> client message
 	s.ToSend <- sent
 	require.NoError(t, conn.SetReadDeadline(time.Now().Add(20*time.Millisecond)))
@@ -99,15 +34,20 @@ func TestMockServer(t *testing.T) {
 }
 
 func TestWSinitMsg(t *testing.T) {
-	s := newMockServer()
+	s := newMockServer(0)
 	defer s.Close()
 	message := "test message"
 	options.initMsg = message
-	options.authHeader += "Bearer ajshdkjhipuqofqldbclqwehqlieh;#kqnwe;ldk"
-	defer func() { options.initMsg = "" }()
+	options.authHeader = "Bearer the_token_is_here"
+	defer func() {
+		options.initMsg = ""
+		options.authHeader = ""
+	}()
 	cmd := &cobra.Command{}
 	time.AfterFunc(100*time.Millisecond, func() { session.cancel() })
 	root(cmd, []string{mockURL})
+	require.Eventually(t, func() bool { return len(s.Received) > 0 }, 20*time.Millisecond, 2*time.Millisecond)
+	require.Equal(t, message, <-s.Received)
 }
 
 func TestWSconnectFail(t *testing.T) {