From faeb37ed80e1f2ae2bfa0c096b2a660c271adddb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Hoguin?= <essen@ninenines.eu>
Date: Sat, 5 Jan 2013 23:35:30 +0100
Subject: Add cowboy_req:set_resp_body_fun/2

This allows streaming a body without knowing the length in advance.
Also allows {stream, StreamFun} response body in the REST code.
---
 test/http_SUITE.erl               | 22 ++++++++++++++++++++++
 test/http_handler_stream_body.erl | 11 ++++++++---
 2 files changed, 30 insertions(+), 3 deletions(-)

(limited to 'test')

diff --git a/test/http_SUITE.erl b/test/http_SUITE.erl
index cd29c87..7ce0835 100644
--- a/test/http_SUITE.erl
+++ b/test/http_SUITE.erl
@@ -68,6 +68,7 @@
 -export([static_test_file/1]).
 -export([static_test_file_css/1]).
 -export([stream_body_set_resp/1]).
+-export([stream_body_set_resp_close/1]).
 -export([te_chunked/1]).
 -export([te_chunked_delayed/1]).
 -export([te_identity/1]).
@@ -117,6 +118,7 @@ groups() ->
 		static_test_file,
 		static_test_file_css,
 		stream_body_set_resp,
+		stream_body_set_resp_close,
 		te_chunked,
 		te_chunked_delayed,
 		te_identity
@@ -235,6 +237,10 @@ init_dispatch(Config) ->
 				[{body, <<"A flameless dance does not equal a cycle">>}]},
 			{[<<"stream_body">>, <<"set_resp">>], http_handler_stream_body,
 				[{reply, set_resp}, {body, <<"stream_body_set_resp">>}]},
+			{[<<"stream_body">>, <<"set_resp_close">>],
+				http_handler_stream_body, [
+					{reply, set_resp_close},
+					{body, <<"stream_body_set_resp_close">>}]},
 			{[<<"static">>, '...'], cowboy_static,
 				[{directory, ?config(static_dir, Config)},
 				 {mimetypes, [{<<".css">>, [<<"text/css">>]}]}]},
@@ -892,6 +898,22 @@ stream_body_set_resp(Config) ->
 	{ok, <<"stream_body_set_resp">>, _}
 		= cowboy_client:response_body(Client3).
 
+stream_body_set_resp_close(Config) ->
+	Client = ?config(client, Config),
+	{ok, Client2} = cowboy_client:request(<<"GET">>,
+		build_url("/stream_body/set_resp_close", Config), Client),
+	{ok, 200, _, Client3} = cowboy_client:response(Client2),
+	{ok, Transport, Socket} = cowboy_client:transport(Client3),
+	case element(7, Client3) of
+		<<"stream_body_set_resp_close">> ->
+			ok;
+		Buffer ->
+			{ok, Rest} = Transport:recv(Socket, 26 - size(Buffer), 1000),
+			<<"stream_body_set_resp_close">> = << Buffer/binary, Rest/binary >>,
+			ok
+	end,
+	{error, closed} = Transport:recv(Socket, 0, 1000).
+
 te_chunked(Config) ->
 	Client = ?config(client, Config),
 	Body = list_to_binary(io_lib:format("~p", [lists:seq(1, 100)])),
diff --git a/test/http_handler_stream_body.erl b/test/http_handler_stream_body.erl
index 15870a8..d8f7d91 100644
--- a/test/http_handler_stream_body.erl
+++ b/test/http_handler_stream_body.erl
@@ -12,10 +12,15 @@ init({_Transport, http}, Req, Opts) ->
 	Reply = proplists:get_value(reply, Opts),
 	{ok, Req, #state{headers=Headers, body=Body, reply=Reply}}.
 
-handle(Req, State=#state{headers=_Headers, body=Body, reply=set_resp}) ->
+handle(Req, State=#state{headers=_Headers, body=Body, reply=Reply}) ->
 	SFun = fun(Socket, Transport) -> Transport:send(Socket, Body) end,
-	SLen = iolist_size(Body),
-	Req2 = cowboy_req:set_resp_body_fun(SLen, SFun, Req),
+	Req2 = case Reply of
+		set_resp ->
+			SLen = iolist_size(Body),
+			cowboy_req:set_resp_body_fun(SLen, SFun, Req);
+		set_resp_close ->
+			cowboy_req:set_resp_body_fun(SFun, Req)
+	end,
 	{ok, Req3} = cowboy_req:reply(200, Req2),
 	{ok, Req3, State}.
 
-- 
cgit v1.2.3