send keepalive on init

99designs · Aug 8, 2019 · f93df34 · f93df34
1 parent bcddd7a
commit f93df34
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 0 deletions.
diff --git a/client/websocket.go b/client/websocket.go
@@ -13,6 +13,7 @@ const (
 	connectionInitMsg = "connection_init" // Client -> Server
 	startMsg          = "start"           // Client -> Server
 	connectionAckMsg  = "connection_ack"  // Server -> Client
+	connectionKa      = "ka"              // Server -> Client
 	dataMsg           = "data"            // Server -> Client
 	errorMsg          = "error"           // Server -> Client
 )
@@ -72,10 +73,20 @@ func (p *Client) WebsocketWithPayload(query string, initPayload map[string]inter
 	if err = c.ReadJSON(&ack); err != nil {
 		return errorSubscription(fmt.Errorf("ack: %s", err.Error()))
 	}
+
 	if ack.Type != connectionAckMsg {
 		return errorSubscription(fmt.Errorf("expected ack message, got %#v", ack))
 	}
 
+	var ka operationMessage
+	if err = c.ReadJSON(&ka); err != nil {
+		return errorSubscription(fmt.Errorf("ka: %s", err.Error()))
+	}
+
+	if ka.Type != connectionKa {
+		return errorSubscription(fmt.Errorf("expected ka message, got %#v", ack))
+	}
+
 	if err = c.WriteJSON(operationMessage{Type: startMsg, ID: "1", Payload: requestBody}); err != nil {
 		return errorSubscription(fmt.Errorf("start: %s", err.Error()))
 	}

diff --git a/handler/websocket.go b/handler/websocket.go
@@ -103,6 +103,7 @@ func (c *wsConnection) init() bool {
 		}
 
 		c.write(&operationMessage{Type: connectionAckMsg})
+		c.write(&operationMessage{Type: connectionKeepAliveMsg})
 	case connectionTerminateMsg:
 		c.close(websocket.CloseNormalClosure, "terminated")
 		return false

diff --git a/handler/websocket_test.go b/handler/websocket_test.go
@@ -59,6 +59,7 @@ func TestWebsocket(t *testing.T) {
 		require.NoError(t, c.WriteJSON(&operationMessage{Type: connectionInitMsg}))
 
 		require.Equal(t, connectionAckMsg, readOp(c).Type)
+		require.Equal(t, connectionKeepAliveMsg, readOp(c).Type)
 	})
 
 	t.Run("client can terminate before run", func(t *testing.T) {
@@ -67,6 +68,7 @@ func TestWebsocket(t *testing.T) {
 
 		require.NoError(t, c.WriteJSON(&operationMessage{Type: connectionInitMsg}))
 		require.Equal(t, connectionAckMsg, readOp(c).Type)
+		require.Equal(t, connectionKeepAliveMsg, readOp(c).Type)
 
 		require.NoError(t, c.WriteJSON(&operationMessage{Type: connectionTerminateMsg}))
 
@@ -80,6 +82,7 @@ func TestWebsocket(t *testing.T) {
 
 		require.NoError(t, c.WriteJSON(&operationMessage{Type: connectionInitMsg}))
 		require.Equal(t, connectionAckMsg, readOp(c).Type)
+		require.Equal(t, connectionKeepAliveMsg, readOp(c).Type)
 
 		require.NoError(t, c.WriteJSON(&operationMessage{
 			Type:    startMsg,
@@ -98,6 +101,7 @@ func TestWebsocket(t *testing.T) {
 
 		require.NoError(t, c.WriteJSON(&operationMessage{Type: connectionInitMsg}))
 		require.Equal(t, connectionAckMsg, readOp(c).Type)
+		require.Equal(t, connectionKeepAliveMsg, readOp(c).Type)
 
 		require.NoError(t, c.WriteJSON(&operationMessage{
 			Type:    startMsg,
@@ -138,6 +142,7 @@ func TestWebsocketWithKeepAlive(t *testing.T) {
 
 		require.NoError(t, c.WriteJSON(&operationMessage{Type: connectionInitMsg}))
 		require.Equal(t, connectionAckMsg, readOp(c).Type)
+		require.Equal(t, connectionKeepAliveMsg, readOp(c).Type)
 
 		require.NoError(t, c.WriteJSON(&operationMessage{
 			Type:    startMsg,
@@ -174,6 +179,7 @@ func TestWebsocketInitFunc(t *testing.T) {
 		require.NoError(t, c.WriteJSON(&operationMessage{Type: connectionInitMsg}))
 
 		require.Equal(t, connectionAckMsg, readOp(c).Type)
+		require.Equal(t, connectionKeepAliveMsg, readOp(c).Type)
 	})
 
 	t.Run("accept connection if WebsocketInitFunc is provided and is accepting connection", func(t *testing.T) {
@@ -189,6 +195,7 @@ func TestWebsocketInitFunc(t *testing.T) {
 		require.NoError(t, c.WriteJSON(&operationMessage{Type: connectionInitMsg}))
 
 		require.Equal(t, connectionAckMsg, readOp(c).Type)
+		require.Equal(t, connectionKeepAliveMsg, readOp(c).Type)
 	})
 
 	t.Run("reject connection if WebsocketInitFunc is provided and is accepting connection", func(t *testing.T) {