Added websocket related client config, fixed tests (#5)

This commit adds several websocket client configuration params, namely:
- WithWriteBufferSize
- WithReadBufferSize
- WithHandshakeTimeout
- WithCompression

It adds relevant tests, fixes previous broken tests, some typos and warnings.

Author: dany74q

configuration.go

@@ -110,3 +110,34 @@ func WithRequestTimeout(interval time.Duration) ClientConfiguration {
 		c.requestTimeout = interval
 	}
 }
+
+// WithWriteBufferSize sets the max write buffer size
+// for the websocket frame
+func WithWriteBufferSize(writeBufferSize int) ClientConfiguration {
+	return func(c *Client) {
+		c.conn.SetWriteBufferSize(writeBufferSize)
+	}
+}
+
+// WithReadBufferSize sets the max read buffer size
+// for the websocket frame
+func WithReadBufferSize(readBufferSize int) ClientConfiguration {
+	return func(c *Client) {
+		c.conn.SetReadBufferSize(readBufferSize)
+	}
+}
+
+// WithHandshakeTimeout sets the websocket handshake timeout
+func WithHandshakeTimeout(handshakeTimeout time.Duration) ClientConfiguration {
+	return func(c *Client) {
+		c.conn.SetHandshakeTimeout(handshakeTimeout)
+	}
+}
+
+// WithCompression sets the compression
+// flag for websocket connections
+func WithCompression(enableCompression bool) ClientConfiguration {
+	return func(c *Client) {
+		c.conn.SetCompression(enableCompression)
+	}
+}

gremconnect/dialer.go

@@ -50,6 +50,10 @@ type Dialer interface {
 	SetPingInterval(interval time.Duration)
 	SetWritingWait(interval time.Duration)
 	SetReadingWait(interval time.Duration)
+	SetWriteBufferSize(writeBufferSize int)
+	SetReadBufferSize(readBufferSize int)
+	SetHandshakeTimeout(handshakeTimeout time.Duration)
+	SetCompression(enableCompression bool)
 }
 
 // NewWebSocketDialer returns a new WebSocket dialer to use when
@@ -58,12 +62,16 @@ type Dialer interface {
 // if they're not assigned by DialerConfig functions.
 func NewWebSocketDialer(address string) Dialer {
 	return &WebSocket{
-		timeout:      5 * time.Second,
-		pingInterval: 60 * time.Second,
-		writingWait:  15 * time.Second,
-		readingWait:  15 * time.Second,
-		connected:    false,
-		address:      address,
-		Quit:         make(chan struct{}),
+		timeout:           5 * time.Second,
+		pingInterval:      60 * time.Second,
+		writingWait:       15 * time.Second,
+		readingWait:       15 * time.Second,
+		handshakeTimeout:  5 * time.Second,
+		writeBufferSize:   50 * 1024,
+		readBufferSize:    50 * 1024,
+		enableCompression: false,
+		connected:         false,
+		address:           address,
+		Quit:              make(chan struct{}),
 	}
 }

gremconnect/websocket.go

@@ -34,16 +34,20 @@ import (
 // to dial to the gremlin server and sustain
 // a stable connection by pinging it regularly.
 type WebSocket struct {
-	address      string
-	conn         *websocket.Conn
-	auth         *Auth
-	disposed     bool
-	connected    bool
-	pingInterval time.Duration
-	writingWait  time.Duration
-	readingWait  time.Duration
-	timeout      time.Duration
-	Quit         chan struct{}
+	address           string
+	conn              *websocket.Conn
+	auth              *Auth
+	disposed          bool
+	connected         bool
+	enableCompression bool
+	pingInterval      time.Duration
+	writingWait       time.Duration
+	readingWait       time.Duration
+	timeout           time.Duration
+	handshakeTimeout  time.Duration
+	writeBufferSize   int
+	readBufferSize    int
+	Quit              chan struct{}
 
 	sync.RWMutex
 }
@@ -54,9 +58,10 @@ type WebSocket struct {
 func (ws *WebSocket) Connect() error {
 	var err error
 	dialer := websocket.Dialer{
-		WriteBufferSize:  1024 * 50, // Set up for large messages.
-		ReadBufferSize:   1024 * 50, // Set up for large messages.
-		HandshakeTimeout: 5 * time.Second,
+		WriteBufferSize:   ws.writeBufferSize,
+		ReadBufferSize:    ws.readBufferSize,
+		HandshakeTimeout:  ws.handshakeTimeout,
+		EnableCompression: ws.enableCompression,
 	}
 
 	// Check if the host address already has the proper
@@ -174,7 +179,7 @@ func (ws *WebSocket) Ping(errs chan error) {
 	}
 }
 
-// Configration functions
+// Configuration functions
 
 // SetAuth will set the authentication to this user and pass
 func (ws *WebSocket) SetAuth(user, pass string) {
@@ -200,3 +205,19 @@ func (ws *WebSocket) SetWritingWait(interval time.Duration) {
 func (ws *WebSocket) SetReadingWait(interval time.Duration) {
 	ws.readingWait = interval
 }
+
+func (ws *WebSocket) SetWriteBufferSize(writeBufferSize int) {
+	ws.writeBufferSize = writeBufferSize
+}
+
+func (ws *WebSocket) SetReadBufferSize(readBufferSize int) {
+	ws.readBufferSize = readBufferSize
+}
+
+func (ws *WebSocket) SetHandshakeTimeout(handshakeTimeout time.Duration) {
+	ws.handshakeTimeout = handshakeTimeout
+}
+
+func (ws *WebSocket) SetCompression(enableCompression bool) {
+	ws.enableCompression = enableCompression
+}

gremconnect/websocket_test.go

@@ -343,3 +343,55 @@ func TestSetReadingWait(t *testing.T) {
 		})
 	})
 }
+
+func TestSetWriteBufferSize(t *testing.T) {
+	Convey("Given a WebSocket and a write buffer size", t, func() {
+		dialer := &WebSocket{}
+		writeBufferSize := 512 * 1024
+		Convey("And SetWriteBufferSize is called", func() {
+			dialer.SetWriteBufferSize(writeBufferSize)
+			Convey("Then the write buffer size should be set in the dialer", func() {
+				So(dialer.writeBufferSize, ShouldEqual, writeBufferSize)
+			})
+		})
+	})
+}
+
+func TestSetReadBufferSize(t *testing.T) {
+	Convey("Given a WebSocket and a read buffer size", t, func() {
+		dialer := &WebSocket{}
+		readBufferSize := 256 * 1024
+		Convey("And SetReadBufferSize is called", func() {
+			dialer.SetReadBufferSize(readBufferSize)
+			Convey("Then the read buffer size should be set in the dialer", func() {
+				So(dialer.readBufferSize, ShouldEqual, readBufferSize)
+			})
+		})
+	})
+}
+
+func TestHandshakeTimeout(t *testing.T) {
+	Convey("Given a WebSocket and a handshake timeout", t, func() {
+		dialer := &WebSocket{}
+		handshakeTimeout := time.Second
+		Convey("And SetHandshakeTimeout is called", func() {
+			dialer.SetHandshakeTimeout(handshakeTimeout)
+			Convey("Then the handshake timeout should be set in the dialer", func() {
+				So(dialer.handshakeTimeout, ShouldEqual, handshakeTimeout)
+			})
+		})
+	})
+}
+
+func TestCompression(t *testing.T) {
+	Convey("Given a WebSocket and a compression flag", t, func() {
+		dialer := &WebSocket{}
+		enableCompression := true
+		Convey("And SetCompression is called", func() {
+			dialer.SetCompression(enableCompression)
+			Convey("Then the compression flag should be set in the dialer", func() {
+				So(dialer.enableCompression, ShouldEqual, enableCompression)
+			})
+		})
+	})
+}

manager/query_test.go

@@ -24,30 +24,35 @@ import (
 	"testing"
 	"time"
 
+	. "github.com/smartystreets/goconvey/convey"
+
 	"github.com/northwesternmutual/grammes/gremconnect"
 	"github.com/northwesternmutual/grammes/logging"
-	. "github.com/smartystreets/goconvey/convey"
 )
 
 // MOCKDIALER
 
 type mockDialer gremconnect.WebSocket
 
-func (*mockDialer) Connect() error                   { return connect }
-func (*mockDialer) Close() error                     { return nil }
-func (*mockDialer) Write([]byte) error               { return nil }
-func (m *mockDialer) Read() ([]byte, error)          { return []byte(response), nil }
-func (*mockDialer) Ping(chan error)                  {}
-func (*mockDialer) IsConnected() bool                { return isConnected }
-func (*mockDialer) IsDisposed() bool                 { return isDisposed }
-func (*mockDialer) Auth() (*gremconnect.Auth, error) { return &gremconnect.Auth{}, nil }
-func (*mockDialer) Address() string                  { return "" }
-func (m *mockDialer) GetQuit() chan struct{}         { return make(chan struct{}) }
-func (*mockDialer) SetAuth(string, string)           {}
-func (*mockDialer) SetTimeout(time.Duration)         {}
-func (*mockDialer) SetPingInterval(time.Duration)    {}
-func (*mockDialer) SetWritingWait(time.Duration)     {}
-func (*mockDialer) SetReadingWait(time.Duration)     {}
+func (*mockDialer) Connect() error                    { return connect }
+func (*mockDialer) Close() error                      { return nil }
+func (*mockDialer) Write([]byte) error                { return nil }
+func (m *mockDialer) Read() ([]byte, error)           { return []byte(response), nil }
+func (*mockDialer) Ping(chan error)                   {}
+func (*mockDialer) IsConnected() bool                 { return isConnected }
+func (*mockDialer) IsDisposed() bool                  { return isDisposed }
+func (*mockDialer) Auth() (*gremconnect.Auth, error)  { return &gremconnect.Auth{}, nil }
+func (*mockDialer) Address() string                   { return "" }
+func (m *mockDialer) GetQuit() chan struct{}          { return make(chan struct{}) }
+func (*mockDialer) SetAuth(string, string)            {}
+func (*mockDialer) SetTimeout(time.Duration)          {}
+func (*mockDialer) SetPingInterval(time.Duration)     {}
+func (*mockDialer) SetWritingWait(time.Duration)      {}
+func (*mockDialer) SetReadingWait(time.Duration)      {}
+func (*mockDialer) SetWriteBufferSize(int)            {}
+func (*mockDialer) SetReadBufferSize(int)             {}
+func (*mockDialer) SetHandshakeTimeout(time.Duration) {}
+func (*mockDialer) SetCompression(bool)               {}
 
 func TestSetLoggerQM(t *testing.T) {
 	Convey("Given a dialer, string executor and query manager", t, func() {

query/traversal/addproperty_test.go

@@ -23,8 +23,9 @@ package traversal
 import (
 	"testing"
 
-	"github.com/northwesternmutual/grammes/query/cardinality"
 	. "github.com/smartystreets/goconvey/convey"
+
+	"github.com/northwesternmutual/grammes/query/cardinality"
 )
 
 func TestProperty(t *testing.T) {
@@ -45,7 +46,7 @@ func TestProperty(t *testing.T) {
 		Convey("When 'Property' is called with object strings and ints", func() {
 			result := g.Property(3, 4, 4)
 			Convey("Then result should equal 'g.property(3,4,4)'", func() {
-				So(result.String(), ShouldEqual, "g.property(3,4,4)")
+				So(result.String(), ShouldEqual, "g.property(3l,4l,4l)")
 			})
 		})
 

query/traversal/pagerank_test.go

@@ -39,7 +39,7 @@ func TestPageRank(t *testing.T) {
 		Convey("When 'PageRank' is called with one argument", func() {
 			result := g.PageRank(1.234)
 			Convey("Then result should equal 'g.pageRank(1.234)'", func() {
-				So(result.String(), ShouldEqual, "g.pageRank(1.234)")
+				So(result.String(), ShouldEqual, "g.pageRank(1.234f)")
 			})
 		})
 

query/traversal/skip_test.go

@@ -23,8 +23,9 @@ package traversal
 import (
 	"testing"
 
-	"github.com/northwesternmutual/grammes/query/scope"
 	. "github.com/smartystreets/goconvey/convey"
+
+	"github.com/northwesternmutual/grammes/query/scope"
 )
 
 func TestSkip(t *testing.T) {
@@ -40,7 +41,7 @@ func TestSkip(t *testing.T) {
 		Convey("When 'Skip' is called with multiple extraFloat arguments", func() {
 			result := g.Skip(scope.Local, 1.234)
 			Convey("Then result should equal 'g.skip(local,1.234)'", func() {
-				So(result.String(), ShouldEqual, "g.skip(local,1.234)")
+				So(result.String(), ShouldEqual, "g.skip(local,1.234f)")
 			})
 		})
 	})

query/traversal/timeLimit_test.go

@@ -32,7 +32,7 @@ func TestTimeLimit(t *testing.T) {
 		Convey("When 'Store' is called", func() {
 			result := g.TimeLimit(1.234)
 			Convey("Then result should equal 'g.timeLimit(1.234)'", func() {
-				So(result.String(), ShouldEqual, "g.timeLimit(1.234)")
+				So(result.String(), ShouldEqual, "g.timeLimit(1.234f)")
 			})
 		})
 	})

query/traversal/util_test.go

@@ -65,7 +65,7 @@ func TestAddStep(t *testing.T) {
 			i := float64(1.234)
 			g.AddStep("test", i)
 			Convey("Then g should equal g.test(1.234)", func() {
-				So(g.String(), ShouldEqual, "g.test(1.234)")
+				So(g.String(), ShouldEqual, "g.test(1.234f)")
 			})
 		})
 
@@ -81,7 +81,7 @@ func TestAddStep(t *testing.T) {
 			s := `a 'test' with \ and \\ and nested {"key": "value", "key2": "val\"val"} and {'key': 'value', 'kay2': 'val\'val'}`
 			g.AddStep("test", s)
 			expected := `g.test('a \'test\' with \\ and \\\\ and nested {"key": "value", "key2": "val\\"val"} and {\'key\': \'value\', \'kay2\': \'val\\\'val\'}')`
-			Convey("Then g should equal " + expected, func() {
+			Convey("Then g should equal "+expected, func() {
 				So(g.String(), ShouldEqual, expected)
 			})
 		})