conver wait for response into a command

cmd/wsget/main.go

@@ -3,6 +3,7 @@ package main
 import (
 	"fmt"
 	"os"
+	"time"
 
 	"github.com/fatih/color"
 	"github.com/ksysoev/wsget/pkg/cli"
@@ -75,7 +76,7 @@ func run(cmd *cobra.Command, args []string) {
 		return
 	}
 
-	wsConn, err := ws.NewWS(wsURL, ws.Options{SkipSSLVerification: insecure, Headers: headers, WaitForResp: waitResponse})
+	wsConn, err := ws.NewWS(wsURL, ws.Options{SkipSSLVerification: insecure, Headers: headers})
 	if err != nil {
 		color.New(color.FgRed).Println("Unable to connect to the server: ", err)
 		return
@@ -101,6 +102,14 @@ func run(cmd *cobra.Command, args []string) {
 
 	if request != "" {
 		opts.Commands = []cli.Executer{cli.NewCommandSend(request)}
+
+		if waitResponse >= 0 {
+			opts.Commands = append(
+				opts.Commands,
+				cli.NewCommandWaitForResp(time.Duration(waitResponse)*time.Second),
+				cli.NewCommandExit(),
+			)
+		}
 	} else {
 		opts.Commands = []cli.Executer{cli.NewCommandEdit("")}
 	}

pkg/cli/cli_test.go

@@ -54,7 +54,7 @@ func TestNewCLI(t *testing.T) {
 		t.Error("Expected non-nil editor")
 	}
 
-	if err = wsConn.Send("Hello, world!"); err != nil {
+	if _, err = wsConn.Send("Hello, world!"); err != nil {
 		t.Fatalf("Unexpected error: %v", err)
 	}
 

pkg/cli/commands.go

@@ -3,6 +3,7 @@ package cli
 import (
 	"fmt"
 	"io"
+	"time"
 
 	"github.com/eiannone/keyboard"
 	"github.com/fatih/color"
@@ -55,11 +56,12 @@ func NewCommandSend(request string) *CommandSend {
 }
 
 func (c *CommandSend) Execute(exCtx *ExecutionContext) (Executer, error) {
-	if err := exCtx.wsConn.Send(c.request); err != nil {
+	msg, err := exCtx.wsConn.Send(c.request)
+	if err != nil {
 		return nil, fmt.Errorf("fail to send request: %s", err)
 	}
 
-	return nil, nil
+	return NewCommandPrintMsg(*msg), nil
 }
 
 type CommandPrintMsg struct {
@@ -110,3 +112,33 @@ func NewCommandExit() *CommandExit {
 func (c *CommandExit) Execute(_ *ExecutionContext) (Executer, error) {
 	return nil, fmt.Errorf("interrupted")
 }
+
+type CommandWaitForResp struct {
+	timeout time.Duration
+}
+
+func NewCommandWaitForResp(timeout time.Duration) *CommandWaitForResp {
+	return &CommandWaitForResp{timeout}
+}
+
+func (c *CommandWaitForResp) Execute(exCtx *ExecutionContext) (Executer, error) {
+	if c.timeout.Seconds() == 0 {
+		msg, ok := <-exCtx.wsConn.Messages
+		if !ok {
+			return nil, fmt.Errorf("connection closed")
+		}
+
+		return NewCommandPrintMsg(msg), nil
+	}
+
+	select {
+	case <-time.After(c.timeout):
+		return nil, fmt.Errorf("timeout")
+	case msg, ok := <-exCtx.wsConn.Messages:
+		if !ok {
+			return nil, fmt.Errorf("connection closed")
+		}
+
+		return NewCommandPrintMsg(msg), nil
+	}
+}

pkg/ws/ws.go

@@ -6,7 +6,6 @@ import (
 	"net/http"
 	"strings"
 	"sync"
-	"time"
 
 	"github.com/fatih/color"
 	"golang.org/x/net/websocket"
@@ -50,7 +49,6 @@ type Connection struct {
 type Options struct {
 	Headers             []string
 	SkipSSLVerification bool
-	WaitForResp         int
 }
 
 func NewWS(url string, opts Options) (*Connection, error) {
@@ -86,14 +84,6 @@ func NewWS(url string, opts Options) (*Connection, error) {
 
 	ws, err := websocket.DialConfig(cfg)
 
-	if opts.WaitForResp > 0 {
-		go func() {
-			time.Sleep(time.Duration(opts.WaitForResp) * time.Second)
-			color.New(color.FgRed).Println("Timeout reached. Closing connection")
-			ws.Close()
-		}()
-	}
-
 	if err != nil {
 		return nil, err
 	}
@@ -113,12 +103,6 @@ func NewWS(url string, opts Options) (*Connection, error) {
 
 			err = websocket.Message.Receive(ws, &msg)
 			if err != nil {
-				if opts.WaitForResp >= 0 {
-					// If we are waiting for single response and connection is closed
-					// we just return from the function
-					return
-				}
-
 				if err.Error() == "EOF" {
 					color.New(color.FgRed).Println("Connection closed by the server")
 				} else {
@@ -129,29 +113,23 @@ func NewWS(url string, opts Options) (*Connection, error) {
 			}
 
 			messages <- Message{Type: Response, Data: msg}
-
-			if opts.WaitForResp >= 0 {
-				return
-			}
 		}
 	}()
 
 	return &Connection{ws: ws, Messages: messages, waitGroup: &waitGroup}, nil
 }
 
-func (wsInsp *Connection) Send(msg string) error {
+func (wsInsp *Connection) Send(msg string) (*Message, error) {
 	wsInsp.waitGroup.Add(1)
 	defer wsInsp.waitGroup.Done()
 
 	err := websocket.Message.Send(wsInsp.ws, msg)
 
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	wsInsp.Messages <- Message{Type: Request, Data: msg}
-
-	return nil
+	return &Message{Type: Request, Data: msg}, nil
 }
 
 func (wsInsp *Connection) Close() {

pkg/ws/ws_test.go

@@ -28,17 +28,13 @@ func TestNewWS(t *testing.T) {
 		t.Fatalf("Expected ws connection, but got nil")
 	}
 
-	if err = ws.Send("Hello, world!"); err != nil {
+	msg, err := ws.Send("Hello, world!")
+	if err != nil {
 		t.Fatalf("Unexpected error: %v", err)
 	}
 
-	select {
-	case msg := <-ws.Messages:
-		if msg.Data != "Hello, world!" {
-			t.Errorf("Expected message data to be 'Hello, world!', but got %v", msg.Data)
-		}
-	default:
-		t.Errorf("Expected message, but got none")
+	if msg.Data != "Hello, world!" {
+		t.Errorf("Expected message data to be 'Hello, world!', but got %v", msg.Data)
 	}
 }