Add the rate_limited/2 REST callback

2 files changed, 50 insertions, 0 deletions

diff --git a/test/handlers/rate_limited_h.erl b/test/handlers/rate_limited_h.erl
new file mode 100644
index 0000000..e54249c
--- /dev/null
+++ b/test/handlers/rate_limited_h.erl
@@ -0,0 +1,24 @@
+%% This module does rate limiting based on the query string value.
+
+-module(rate_limited_h).
+
+-export([init/2]).
+-export([rate_limited/2]).
+-export([content_types_provided/2]).
+-export([get_text_plain/2]).
+
+init(Req, State) ->
+	{cowboy_rest, Req, State}.
+
+rate_limited(Req=#{qs := <<"false">>}, State) ->
+	{false, Req, State};
+rate_limited(Req=#{qs := <<"true-date">>}, State) ->
+	{{true, {{2222, 2, 22}, {11, 11, 11}}}, Req, State};
+rate_limited(Req=#{qs := <<"true">>}, State) ->
+	{{true, 3600}, Req, State}.
+
+content_types_provided(Req, State) ->
+	{[{{<<"text">>, <<"plain">>, []}, get_text_plain}], Req, State}.
+
+get_text_plain(Req, State) ->
+	{<<"This is REST!">>, Req, State}.
diff --git a/test/rest_handler_SUITE.erl b/test/rest_handler_SUITE.erl
index 93532fe..09ce8fb 100644
--- a/test/rest_handler_SUITE.erl
+++ b/test/rest_handler_SUITE.erl
@@ -48,6 +48,7 @@ init_dispatch(_) ->
 		{"/charset_in_content_types_provided_implicit_no_callback",
 			charset_in_content_types_provided_implicit_no_callback_h, []},
 		{"/provide_callback_missing", provide_callback_missing_h, []},
+		{"/rate_limited", rate_limited_h, []},
 		{"/switch_handler", switch_handler_h, run},
 		{"/switch_handler_opts", switch_handler_h, hibernate}
 	]}]).
@@ -287,6 +288,31 @@ provide_callback_missing(Config) ->
 	{response, fin, 500, _} = gun:await(ConnPid, Ref),
 	ok.
 
+rate_limited(Config) ->
+	doc("A 429 response must be sent when the rate_limited callback returns true. "
+		"The retry-after header is specified as an integer."),
+	ConnPid = gun_open(Config),
+	Ref = gun:get(ConnPid, "/rate_limited?true", [{<<"accept-encoding">>, <<"gzip">>}]),
+	{response, fin, 429, Headers} = gun:await(ConnPid, Ref),
+	{_, <<"3600">>} = lists:keyfind(<<"retry-after">>, 1, Headers),
+	ok.
+
+rate_limited_datetime(Config) ->
+	doc("A 429 response must be sent when the rate_limited callback returns true. "
+		"The retry-after header is specified as a date/time tuple."),
+	ConnPid = gun_open(Config),
+	Ref = gun:get(ConnPid, "/rate_limited?true-date", [{<<"accept-encoding">>, <<"gzip">>}]),
+	{response, fin, 429, Headers} = gun:await(ConnPid, Ref),
+	{_, <<"Fri, 22 Feb 2222 11:11:11 GMT">>} = lists:keyfind(<<"retry-after">>, 1, Headers),
+	ok.
+
+rate_not_limited(Config) ->
+	doc("A success response must be sent when the rate_limited callback returns false."),
+	ConnPid = gun_open(Config),
+	Ref = gun:get(ConnPid, "/rate_limited?false", [{<<"accept-encoding">>, <<"gzip">>}]),
+	{response, nofin, 200, _} = gun:await(ConnPid, Ref),
+	ok.
+
 switch_handler(Config) ->
 	doc("Switch REST to loop handler for streaming the response body."),
 	ConnPid = gun_open(Config),