optimize echo server and tests

connection_test.go

@@ -14,6 +14,8 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func success() error { return nil }
+
 func TestGetPrefix(t *testing.T) {
 	options.timestamp = true
 	defer func() { options.timestamp = false }()
@@ -31,7 +33,9 @@ func TestSession(t *testing.T) {
 	conn := newMockConn()
 	defer TryCloseNormally(conn, "test finished")
 	outR, outW, _ := os.Pipe()
-	rl, err := readline.NewEx(&readline.Config{Prompt: "> ", Stdout: outW})
+	// substitute FuncMakeRaw and FuncExitRaw to empty func to use open pipe as Stdin
+	// switching to raw file descriptor 0 will cause immediately closure of rl due to EOF
+	rl, err := readline.NewEx(&readline.Config{Prompt: "> ", Stdout: outW, FuncMakeRaw: success, FuncExitRaw: success})
 	require.NoError(t, err)
 	s := Session{
 		ws:      conn,
@@ -42,7 +46,7 @@ func TestSession(t *testing.T) {
 	}
 	sent := "test message"
 	typed := "typed"
-	binary := "binary"
+	binary := "binary as text"
 	unknown := "unknown"
 	toBeFiltered := "must be filtered"
 	// test sendMsg
@@ -60,6 +64,7 @@ func TestSession(t *testing.T) {
 	}()
 	require.Eventually(t, func() bool { return len(srv.Received) > 0 }, 20*time.Millisecond, 2*time.Millisecond)
 	require.Equal(t, typed, <-srv.Received)
+	require.NoError(t, rl.Close())
 	// test readWebsocket
 	go func() {
 		s.readWebsocket()
@@ -89,16 +94,16 @@ func TestSession(t *testing.T) {
 	time.Sleep(20 * time.Millisecond)
 	outW.Close()
 	output, err := io.ReadAll(outR)
-	out := string(output)
 	require.NoError(t, err)
+	out := string(output)
 	require.Contains(t, out, " > "+sent)
 	require.Contains(t, out, " > "+typed)
 	require.Contains(t, out, " < "+sent)
 	require.Contains(t, out, " < \n00000000  74 65 73 74 20 6d 65 73  73 61 67 65              |"+sent+"|")
 	require.Contains(t, out, " < "+binary)
 	require.NotContains(t, out, toBeFiltered)
 	require.NotContains(t, out, unknown)
-	// t.Log(out)
+	t.Log(out)
 }
 
 func TestPingPong(t *testing.T) {
@@ -117,14 +122,13 @@ func TestPingPong(t *testing.T) {
 	errs := make(chan []error, 1)
 	// substitute FuncMakeRaw and FuncExitRaw to empty func to use open pipe as Stdin
 	// switching to raw file descriptor 0 will cause immediately closure of rl due to EOF
-	success := func() error { return nil }
 	rl, err := readline.NewEx(&readline.Config{Prompt: "> ", Stdin: inR, Stdout: outW, FuncMakeRaw: success, FuncExitRaw: success})
 	require.NoError(t, err)
 	s := &Session{rl: rl}
 	go func() {
 		errs <- s.connect(mockURL)
 	}()
-	time.Sleep(20 * time.Millisecond)
+	time.Sleep(20 * time.Millisecond) // let it work for a while
 	require.Eventually(t, func() bool { return s.cancel != nil }, 100*time.Millisecond, 2*time.Millisecond)
 	s.cancel()
 	inW.Close()
@@ -151,9 +155,12 @@ func TestInitMsg(t *testing.T) {
 	rl, err := readline.New(" >")
 	require.NoError(t, err)
 	s := &Session{rl: rl}
-	time.AfterFunc(500*time.Millisecond, func() { s.cancel() })
-	errs := s.connect(mockURL)
-	require.Empty(t, errs)
+	errs := make(chan []error)
+	go func() {
+		errs <- s.connect(mockURL)
+	}()
 	require.Eventually(t, func() bool { return len(m.Received) > 0 }, 20*time.Millisecond, 2*time.Millisecond)
 	require.Equal(t, message, <-m.Received)
+	s.cancel()
+	require.Empty(t, <-errs)
 }

echo-server/echo-server.go

@@ -23,29 +23,16 @@ type msg struct {
 	ListenerCount *int   `json:"listenerCount,omitempty"`
 }
 
-func sendMsg(m msg, c *websocket.Conn) {
-	response, _ := json.Marshal(m)
-	c.WriteMessage(websocket.TextMessage, response)
-}
-
-func main() {
-	DoMain(os.Args)
+type echoServer struct {
+	*server.Server
 }
 
-var srv *server.Server
-
-// DoMain is main with os.Args as parameter
-func DoMain(args []string) {
-	raw := defaultUrl
-	if len(args) > 1 {
-		raw = args[1]
-	}
-	u, err := url.Parse(raw)
+func New(addr string) (*echoServer, error) {
+	u, err := url.Parse(addr)
 	if err != nil {
-		fmt.Printf("url parsing error: %v", err)
-		os.Exit(1)
+		return nil, fmt.Errorf("url parsing error: %v", err)
 	}
-	srv = server.NewServer(u.Host)
+	srv := server.NewServer(u.Host)
 	if !strings.HasPrefix(u.Path, "/") {
 		u.Path = "/" + u.Path
 	}
@@ -83,16 +70,41 @@ func DoMain(args []string) {
 			}
 		}
 	})
+	return &echoServer{srv}, nil
+}
+
+func sendMsg(m msg, c *websocket.Conn) {
+	response, _ := json.Marshal(m)
+	c.WriteMessage(websocket.TextMessage, response)
+}
+
+func main() {
+	if err := DoMain(os.Args); err != nil {
+		fmt.Printf("%v\n", err)
+		os.Exit(1)
+	}
+}
+
+// DoMain is main with os.Args as parameter
+func DoMain(args []string) error {
+	addr := defaultUrl
+	if len(args) > 1 {
+		addr = args[1]
+	}
+	srv, err := New(addr)
+	if err != nil {
+		return err
+	}
 	go func() {
 		sig := make(chan os.Signal, 2)
 		signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
 		fmt.Printf("\n%s signal received, exiting...\n", <-sig)
 		srv.Close()
 	}()
-	fmt.Printf("starting echo server on %s...\n", u.Host)
+	fmt.Printf("starting echo server on %s...\n", addr)
 	err = srv.ListenAndServe()
 	if !errors.Is(err, http.ErrServerClosed) {
-		fmt.Println(err)
-		os.Exit(1)
+		return err
 	}
+	return nil
 }

echo-server/echo-server_test.go

@@ -16,32 +16,12 @@ import (
 )
 
 func TestEchoServer(t *testing.T) {
-	envName := fmt.Sprintf("BE_%s", t.Name())
-	if os.Getenv(envName) == "1" {
-		// conn, _, err := dialer.Dial(defaultUrl, nil)
-		// if err == nil {
-		// 	fmt.Println("unexpected connection")
-		// 	os.Exit(2)
-		// }
-		go DoMain([]string{"", defaultUrl})
-		time.Sleep(500 * time.Millisecond)
-		// server.TryCloseNormally(conn, "tests finish")
-		//require.NoError(t, srv.Close())
-		syscall.Kill(syscall.Getpid(), syscall.SIGINT)
-		return
-	}
-	// collect child coverage artifacts in the parent files
-	args := []string{"-test.run=" + t.Name()}
-	for _, v := range os.Args {
-		if strings.Contains(v, "cover") {
-			args = append(args, v)
-		}
-	}
-	cmd := exec.Command(os.Args[0], args...)
-	r, err := cmd.StdoutPipe()
+	srv, err := New(defaultUrl)
+	errCh := make(chan error)
+	go func() {
+		errCh <- srv.ListenAndServe()
+	}()
 	require.NoError(t, err)
-	cmd.Env = append(os.Environ(), envName+"=1")
-	require.NoError(t, cmd.Start())
 	dialer := websocket.Dialer{}
 	var conn *websocket.Conn
 	require.Eventually(t, func() bool {
@@ -95,18 +75,35 @@ func TestEchoServer(t *testing.T) {
 			}
 		})
 	}
+	err = srv.Close()
+	require.NoError(t, err)
+	require.EqualError(t, <-errCh, "http: Server closed")
+}
 
-	out, _ := io.ReadAll(r)
-	output := string(out)
-	err = cmd.Wait()
+func TestEchoServerConErr(t *testing.T) {
+	srv, err := New(defaultUrl)
+	errCh := make(chan error)
+	go func() {
+		errCh <- srv.ListenAndServe()
+	}()
 	require.NoError(t, err)
-	require.Contains(t, output, "starting echo server on localhost:8080...")
+	dialer := websocket.Dialer{}
+	var conn *websocket.Conn
+	require.Eventually(t, func() bool {
+		conn, _, err = dialer.Dial(defaultUrl, nil)
+		return err == nil
+	}, 50*time.Millisecond, 5*time.Millisecond)
+	require.NoError(t, conn.Close())
+	require.NoError(t, srv.Close())
 }
 
-func TestDoMain(t *testing.T) {
+func TestMainDown(t *testing.T) {
 	envName := fmt.Sprintf("BE_%s", t.Name())
 	if os.Getenv(envName) == "1" {
-		DoMain([]string{"", ":"})
+		os.Args = []string([]string{""})
+		go main()
+		time.Sleep(50 * time.Millisecond)
+		syscall.Kill(syscall.Getpid(), syscall.SIGINT)
 		return
 	}
 	args := []string{"-test.run=" + t.Name()}
@@ -122,25 +119,18 @@ func TestDoMain(t *testing.T) {
 	require.NoError(t, cmd.Start())
 	out, _ := io.ReadAll(r)
 	err = cmd.Wait()
-	require.EqualError(t, err, "exit status 1")
-	require.Equal(t, "url parsing error: parse \":\": missing protocol scheme", string(out))
+	require.NoError(t, err)
+	require.Contains(t, string(out), "starting echo server on ws://localhost:8080/ws...\n")
 }
 
-func TestEchoServer2(t *testing.T) {
-	url := "ws://localhost:8080"
+func TestMainError(t *testing.T) {
 	envName := fmt.Sprintf("BE_%s", t.Name())
 	if os.Getenv(envName) == "1" {
-		go DoMain([]string{"", url})
-		time.Sleep(50 * time.Millisecond)
-		dialer := websocket.Dialer{}
-		conn, _, err := dialer.Dial(url, nil)
-		require.NoError(t, err)
-		conn.Close()
-		time.Sleep(5 * time.Millisecond)
-		syscall.Kill(syscall.Getpid(), syscall.SIGINT)
+		os.Args = []string([]string{"", ":"})
+		main()
+
 		return
 	}
-	// collect child coverage artifacts in the parent files
 	args := []string{"-test.run=" + t.Name()}
 	for _, v := range os.Args {
 		if strings.Contains(v, "cover") {
@@ -153,8 +143,7 @@ func TestEchoServer2(t *testing.T) {
 	cmd.Env = append(os.Environ(), envName+"=1")
 	require.NoError(t, cmd.Start())
 	out, _ := io.ReadAll(r)
-	output := string(out)
 	err = cmd.Wait()
-	require.NoError(t, err)
-	require.Contains(t, output, "starting echo server on localhost:8080...")
+	require.EqualError(t, err, "exit status 1")
+	require.Equal(t, "url parsing error: parse \":\": missing protocol scheme\n", string(out))
 }

main.go

@@ -69,10 +69,14 @@ func root(cmd *cobra.Command, args []string) {
 	}
 	if options.origin == "" {
 		originURL := *dest
-		if dest.Scheme == "wss" {
+		switch dest.Scheme {
+		case "wss":
 			originURL.Scheme = "https"
-		} else {
+		case "ws":
 			originURL.Scheme = "http"
+		default:
+			fmt.Printf("unsupported scheme: %s\n", dest.Scheme)
+			os.Exit(1)
 		}
 		options.origin = originURL.String()
 	}

mockServer.go

@@ -27,8 +27,8 @@ func newMockServer(interval time.Duration) *mockServer {
 	s := server.NewServer(u.Host)
 	m := &mockServer{
 		Close: func() error {
-			cancel()
 			s.Shutdown(ctx)
+			cancel()
 			s.Close()
 			return nil
 		},