From d63eccd8b84d376c288ed7b056f574d34a5ffc4a Mon Sep 17 00:00:00 2001
From: Feng Lee <feng@emqx.io>
Date: Fri, 16 Aug 2019 18:24:39 +0800
Subject: [PATCH] Add test cases for emqx_keepalive module (#2784)

* Improve the keepalive module
---
 src/emqx_keepalive.erl        | 34 ++++++++++++++++++----------------
 test/emqx_keepalive_SUITE.erl |  4 ++++
 2 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/src/emqx_keepalive.erl b/src/emqx_keepalive.erl
index 25170067d..88848f7ac 100644
--- a/src/emqx_keepalive.erl
+++ b/src/emqx_keepalive.erl
@@ -25,14 +25,16 @@
 -export_type([keepalive/0]).
 
 -record(keepalive, {
-          statfun,
-          statval,
-          tsec,
-          tmsg,
-          tref,
-          repeat = 0
+          statfun    :: statfun(),
+          statval    :: integer(),
+          tsec       :: pos_integer(),
+          tmsg       :: term(),
+          tref       :: reference(),
+          repeat = 0 :: non_neg_integer()
          }).
 
+-type(statfun() :: fun(() -> {ok, integer()} | {error, term()})).
+
 -opaque(keepalive() :: #keepalive{}).
 
 %%--------------------------------------------------------------------
@@ -40,15 +42,17 @@
 %%--------------------------------------------------------------------
 
 %% @doc Start a keepalive
--spec(start(fun(), integer(), any()) -> {ok, keepalive()} | {error, term()}).
-start(_, 0, _) ->
-    {ok, #keepalive{}};
-start(StatFun, TimeoutSec, TimeoutMsg) ->
+-spec(start(statfun(), pos_integer(), term())
+      -> {ok, keepalive()} | {error, term()}).
+start(StatFun, TimeoutSec, TimeoutMsg) when TimeoutSec > 0 ->
     try StatFun() of
         {ok, StatVal} ->
-            {ok, #keepalive{statfun = StatFun, statval = StatVal,
-                            tsec = TimeoutSec, tmsg = TimeoutMsg,
-                            tref = timer(TimeoutSec, TimeoutMsg)}};
+            TRef = timer(TimeoutSec, TimeoutMsg),
+            {ok, #keepalive{statfun = StatFun,
+                            statval = StatVal,
+                            tsec = TimeoutSec,
+                            tmsg = TimeoutMsg,
+                            tref = TRef}};
         {error, Error} ->
             {error, Error}
     catch
@@ -82,9 +86,7 @@ resume(KeepAlive = #keepalive{tsec = TimeoutSec, tmsg = TimeoutMsg}) ->
 %% @doc Cancel Keepalive
 -spec(cancel(keepalive()) -> ok).
 cancel(#keepalive{tref = TRef}) when is_reference(TRef) ->
-    catch erlang:cancel_timer(TRef), ok;
-cancel(_) ->
-    ok.
+    catch erlang:cancel_timer(TRef), ok.
 
 timer(Secs, Msg) ->
     erlang:send_after(timer:seconds(Secs), self(), Msg).
diff --git a/test/emqx_keepalive_SUITE.erl b/test/emqx_keepalive_SUITE.erl
index abfd8ee2b..f140913ec 100644
--- a/test/emqx_keepalive_SUITE.erl
+++ b/test/emqx_keepalive_SUITE.erl
@@ -35,3 +35,7 @@ keepalive_recv(KA, Acc) ->
         after 4000 -> Acc
     end.
 
+t_cancel(_) ->
+    {ok, KA} = emqx_keepalive:start(fun() -> {ok, 1} end, 1, {keepalive, timeout}),
+    ok = emqx_keepalive:cancel(KA).
+