From e2d5c268aa9dc7d37197560588fe9088741fb6d1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Hoguin?= <essen@ninenines.eu>
Date: Thu, 11 Aug 2016 11:06:03 +0200
Subject: Fix timer-sent events in the websocket example

---
 examples/websocket/src/ws_handler.erl |  6 +++++-
 src/cowboy_websocket.erl              |  6 ++++++
 test/examples_SUITE.erl               | 21 ++++++++++++++++++---
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/examples/websocket/src/ws_handler.erl b/examples/websocket/src/ws_handler.erl
index eaa9284..69d8d7b 100644
--- a/examples/websocket/src/ws_handler.erl
+++ b/examples/websocket/src/ws_handler.erl
@@ -1,13 +1,17 @@
 -module(ws_handler).
 
 -export([init/2]).
+-export([websocket_init/2]).
 -export([websocket_handle/3]).
 -export([websocket_info/3]).
 
 init(Req, Opts) ->
-	erlang:start_timer(1000, self(), <<"Hello!">>),
 	{cowboy_websocket, Req, Opts}.
 
+websocket_init(Req, State) ->
+	erlang:start_timer(1000, self(), <<"Hello!">>),
+	{ok, Req, State}.
+
 websocket_handle({text, Msg}, Req, State) ->
 	{reply, {text, << "That's what she said! ", Msg/binary >>}, Req, State};
 websocket_handle(_Data, Req, State) ->
diff --git a/src/cowboy_websocket.erl b/src/cowboy_websocket.erl
index 693a920..d932323 100644
--- a/src/cowboy_websocket.erl
+++ b/src/cowboy_websocket.erl
@@ -32,6 +32,12 @@
 	| {module(), Req, any(), timeout()}
 	| {module(), Req, any(), timeout(), hibernate}
 	when Req::cowboy_req:req().
+
+-callback websocket_init(Req, State)
+	-> {ok, Req, State}
+	when Req::cowboy_req:req(), State::any().
+-optional_callbacks([websocket_init/2]).
+
 -callback websocket_handle({text | binary | ping | pong, binary()}, Req, State)
 	-> {ok, Req, State}
 	| {ok, Req, State, hibernate}
diff --git a/test/examples_SUITE.erl b/test/examples_SUITE.erl
index b9cc7ba..25cf7e4 100644
--- a/test/examples_SUITE.erl
+++ b/test/examples_SUITE.erl
@@ -393,12 +393,27 @@ websocket(_) ->
 			Msg1 ->
 				exit({connection_failed, Msg1})
 		end,
+		%% Check that we receive the message sent on timer on init.
+		receive
+			{gun_ws, Pid, {text, <<"Hello!">>}} ->
+				ok
+		after 2000 ->
+			exit(timeout)
+		end,
+		%% Check that we receive subsequent messages sent on timer.
+		receive
+			{gun_ws, Pid, {text, <<"How' you doin'?">>}} ->
+				ok
+		after 2000 ->
+			exit(timeout)
+		end,
+		%% Check that we receive the echoed message.
 		gun:ws_send(Pid, {text, <<"hello">>}),
 		receive
 			{gun_ws, Pid, {text, <<"That's what she said! hello">>}} ->
-				ok;
-			Msg2 ->
-				exit({receive_failed, Msg2})
+				ok
+		after 500 ->
+			exit(timeout)
 		end,
 		gun:ws_send(Pid, close)
 	after
-- 
cgit v1.2.3