Fix test

helium · Sep 26, 2023 · 148e164 · 148e164
1 parent ce4bc23
commit 148e164
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 17 deletions.
diff --git a/test/hpr_protocol_router_SUITE.erl b/test/hpr_protocol_router_SUITE.erl
@@ -320,20 +320,20 @@ gateway_disconnect_test(_Config) ->
     _ConnPid = h2_stream_set:connection(StreamSet),
 
     ok = gen_server:stop(GatewayPid),
-    ok =
-        receive
-            {hpr_test_gateway, GatewayPid,
-                {terminate, #{channel := GatewayStreamSet, stream_pid := GatewayStreamPid}}} ->
-                GatewayConnPid = h2_stream_set:connection(GatewayStreamSet),
-                ok = test_utils:wait_until(
-                    fun() ->
-                        false == erlang:is_process_alive(GatewayConnPid) andalso
-                            false == erlang:is_process_alive(GatewayStreamPid) andalso
-                            false == erlang:is_process_alive(GatewayPid)
-                    end
-                )
-        after timer:seconds(3) -> ct:fail(no_terminate_rcvd)
-        end,
+
+    case hpr_test_gateway:receive_terminate(GatewayPid) of
+        {error, timeout} ->
+            ct:fail(no_terminate_rcvd);
+        {ok, #{channel := GatewayStreamSet, stream_pid := GatewayStreamPid}} ->
+            GatewayConnPid = h2_stream_set:connection(GatewayStreamSet),
+            ok = test_utils:wait_until(
+                fun() ->
+                    false == erlang:is_process_alive(GatewayConnPid) andalso
+                        false == erlang:is_process_alive(GatewayStreamPid) andalso
+                        false == erlang:is_process_alive(GatewayPid)
+                end
+            )
+    end,
 
     ok = test_utils:wait_until(
         fun() ->

diff --git a/test/hpr_test_gateway.erl b/test/hpr_test_gateway.erl
@@ -13,7 +13,8 @@
     receive_env_down/1,
     receive_register/1,
     receive_session_init/2,
-    receive_stream_down/1
+    receive_stream_down/1,
+    receive_terminate/1
 ]).
 
 %% ------------------------------------------------------------------
@@ -115,6 +116,15 @@ receive_stream_down(GatewayPid) ->
         {error, timeout}
     end.
 
+-spec receive_terminate(GatewayPid :: pid()) -> {ok, any()} | {error, timeout}.
+receive_terminate(GatewayPid) ->
+    receive
+        {?MODULE, GatewayPid, {terminate, Stream}} ->
+            {ok, Stream}
+    after timer:seconds(2) ->
+        {error, timeout}
+    end.
+
 %% ------------------------------------------------------------------
 %%% gen_server Function Definitions
 %% ------------------------------------------------------------------
@@ -279,14 +289,16 @@ handle_info(_Msg, State) ->
     lager:debug("unknown info ~p", [_Msg]),
     {noreply, State}.
 
-terminate(_Reason, #state{pubkey_bin = PubKeyBin, stream = undefined}) ->
+terminate(_Reason, #state{forward = Pid, pubkey_bin = PubKeyBin, stream = undefined}) ->
     ok = grpcbox_channel:stop(PubKeyBin),
     lager:debug("terminate ~p", [_Reason]),
+    Pid ! {?MODULE, self(), {terminate, undefined}},
     ok;
-terminate(_Reason, #state{pubkey_bin = PubKeyBin, stream = Stream}) ->
+terminate(_Reason, #state{forward = Pid, pubkey_bin = PubKeyBin, stream = Stream}) ->
     ok = grpcbox_client:close_send(Stream),
     ok = grpcbox_channel:stop(PubKeyBin),
     lager:debug("terminate ~p", [_Reason]),
+    Pid ! {?MODULE, self(), {terminate, Stream}},
     ok.
 
 %% ------------------------------------------------------------------