diff --git a/pkg/client/client.go b/pkg/client/client.go
index 842c937..a43cd52 100644
--- a/pkg/client/client.go
+++ b/pkg/client/client.go
@@ -26,7 +26,10 @@ import (
 	"google.golang.org/grpc/metadata"
 )
 
-const defaultHTTPTimeout = 30 * time.Second
+const (
+	defaultHTTPTimeout = 30 * time.Second
+	waitParameterName  = "wait"
+)
 
 type APIError struct {
 	StatusCode int
@@ -340,30 +343,6 @@ func (client *Client) wsRequestRaw(
 	return conn, nil
 }
 
-func (client *Client) httpClientForWebSocket(params map[string]string) *http.Client {
-	waitRaw, ok := params["wait"]
-	if !ok {
-		return client.httpClient
-	}
-
-	waitSeconds, err := strconv.ParseUint(waitRaw, 10, 16)
-	if err != nil {
-		return client.httpClient
-	}
-
-	waitTimeout := time.Duration(waitSeconds)*time.Second + defaultHTTPTimeout
-	if client.httpClient.Timeout == 0 || client.httpClient.Timeout >= waitTimeout {
-		return client.httpClient
-	}
-
-	return &http.Client{
-		CheckRedirect: client.httpClient.CheckRedirect,
-		Jar:           client.httpClient.Jar,
-		Timeout:       waitTimeout,
-		Transport:     client.httpClient.Transport,
-	}
-}
-
 func (client *Client) formatPath(path string) *url.URL {
 	endpointURL := &url.URL{
 		Scheme: client.baseURL.Scheme,
@@ -424,3 +403,27 @@ func (client *Client) RPC() *RPCService {
 		client: client,
 	}
 }
+
+func (client *Client) httpClientForWebSocket(params map[string]string) *http.Client {
+	waitRaw, ok := params[waitParameterName]
+	if !ok {
+		return client.httpClient
+	}
+
+	waitSeconds, err := strconv.ParseUint(waitRaw, 10, 16)
+	if err != nil {
+		return client.httpClient
+	}
+
+	waitTimeout := time.Duration(waitSeconds)*time.Second + defaultHTTPTimeout
+	if client.httpClient.Timeout == 0 || client.httpClient.Timeout >= waitTimeout {
+		return client.httpClient
+	}
+
+	return &http.Client{
+		CheckRedirect: client.httpClient.CheckRedirect,
+		Jar:           client.httpClient.Jar,
+		Timeout:       waitTimeout,
+		Transport:     client.httpClient.Transport,
+	}
+}
diff --git a/pkg/client/vms_test.go b/pkg/client/vms_test.go
index 8fe2e79..0e4ca23 100644
--- a/pkg/client/vms_test.go
+++ b/pkg/client/vms_test.go
@@ -14,7 +14,7 @@ func TestHTTPClientForWebSocketHonorsWait(t *testing.T) {
 	devClient, err := New(WithAddress("http://localhost"))
 	require.NoError(t, err)
 
-	httpClient := devClient.httpClientForWebSocket(map[string]string{"wait": "120"})
+	httpClient := devClient.httpClientForWebSocket(map[string]string{waitParameterName: "120"})
 
 	require.Equal(t, 150*time.Second, httpClient.Timeout)
 	require.Same(t, devClient.httpClient.Transport, httpClient.Transport)
@@ -28,7 +28,12 @@ func TestExecSessionBuildsReconnectableQuery(t *testing.T) {
 			query = request.URL.Query()
 
 			conn, err := websocket.Accept(writer, request, nil)
-			require.NoError(t, err)
+			if err != nil {
+				t.Errorf("failed to accept WebSocket connection: %v", err)
+
+				return
+			}
+
 			defer conn.CloseNow()
 		}),
 	)
@@ -58,6 +63,6 @@ func TestExecSessionBuildsReconnectableQuery(t *testing.T) {
 	require.Equal(t, []string{"80"}, query["cols"])
 	require.Equal(t, []string{"hello"}, query["env[GREETING]"])
 	require.Equal(t, []string{"/tmp"}, query["workdir"])
-	require.Equal(t, []string{"7"}, query["wait"])
+	require.Equal(t, []string{"7"}, query[waitParameterName])
 	require.Equal(t, []string{"resume-me"}, query["session"])
 }