Fix timer-sent events in the websocket example

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) ->

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}

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