refactor: Clean up code and improve test coverage

client/client.go

@@ -6,7 +6,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"golang.org/x/net/websocket"
 	"net"
 	"net/http"
 	"net/smtp"
@@ -19,6 +18,7 @@ import (
 	"github.com/ishidawataru/sctp"
 	ping "github.com/prometheus-community/pro-bing"
 	"golang.org/x/crypto/ssh"
+	"golang.org/x/net/websocket"
 )
 
 var (
@@ -261,24 +261,25 @@ func Ping(address string, config *Config) (bool, time.Duration) {
 	return true, 0
 }
 
-// Open a websocket connection, write `body` and return a message from the server
-func QueryWebSocket(address string, config *Config, body string) (bool, []byte, error) {
+// QueryWebSocket opens a websocket connection, write `body` and return a message from the server
+func QueryWebSocket(address, body string, config *Config) (bool, []byte, error) {
 	const (
 		Origin             = "http://localhost/"
 		MaximumMessageSize = 1024 // in bytes
 	)
-
 	wsConfig, err := websocket.NewConfig(address, Origin)
 	if err != nil {
 		return false, nil, fmt.Errorf("error configuring websocket connection: %w", err)
 	}
+	if config != nil {
+		wsConfig.Dialer = &net.Dialer{Timeout: config.Timeout}
+	}
 	// Dial URL
 	ws, err := websocket.DialConfig(wsConfig)
 	if err != nil {
 		return false, nil, fmt.Errorf("error dialing websocket: %w", err)
 	}
 	defer ws.Close()
-	connected := true
 	// Write message
 	if _, err := ws.Write([]byte(body)); err != nil {
 		return false, nil, fmt.Errorf("error writing websocket body: %w", err)
@@ -289,7 +290,7 @@ func QueryWebSocket(address string, config *Config, body string) (bool, []byte,
 	if n, err = ws.Read(msg); err != nil {
 		return false, nil, fmt.Errorf("error reading websocket message: %w", err)
 	}
-	return connected, msg[:n], nil
+	return true, msg[:n], nil
 }
 
 // InjectHTTPClient is used to inject a custom HTTP client for testing purposes

client/client_test.go

@@ -253,3 +253,14 @@ func TestHttpClientProvidesOAuth2BearerToken(t *testing.T) {
 		t.Error("exptected `secret-token` as Bearer token in the mocked response header `X-Org-Authorization`, but got", response.Header.Get("X-Org-Authorization"))
 	}
 }
+
+func TestQueryWebSocket(t *testing.T) {
+	_, _, err := QueryWebSocket("", "body", &Config{Timeout: 2 * time.Second})
+	if err == nil {
+		t.Error("expected an error due to the address being invalid")
+	}
+	_, _, err = QueryWebSocket("ws://example.org", "body", &Config{Timeout: 2 * time.Second})
+	if err == nil {
+		t.Error("expected an error due to the target not being websocket-friendly")
+	}
+}

client/config.go

@@ -16,7 +16,7 @@ import (
 )
 
 const (
-	defaultHTTPTimeout = 10 * time.Second
+	defaultTimeout = 10 * time.Second
 )
 
 var (
@@ -27,7 +27,7 @@ var (
 	defaultConfig = Config{
 		Insecure:       false,
 		IgnoreRedirect: false,
-		Timeout:        defaultHTTPTimeout,
+		Timeout:        defaultTimeout,
 	}
 )
 

client/config_test.go

@@ -13,7 +13,7 @@ func TestConfig_getHTTPClient(t *testing.T) {
 	if !(insecureClient.Transport).(*http.Transport).TLSClientConfig.InsecureSkipVerify {
 		t.Error("expected Config.Insecure set to true to cause the HTTP client to skip certificate verification")
 	}
-	if insecureClient.Timeout != defaultHTTPTimeout {
+	if insecureClient.Timeout != defaultTimeout {
 		t.Error("expected Config.Timeout to default the HTTP client to a timeout of 10s")
 	}
 	request, _ := http.NewRequest("GET", "", nil)