test(emqx): make peeking into connection state module-aware

emqx · keynslug · Jul 24, 2025 · Jun 30, 2025 · Jul 23, 2025 · Jun 30, 2025
commit 28bc810e4355d92b5ea78654ac88490037c21b21
@@ -0,0 +1,36 @@
+%%--------------------------------------------------------------------
+%% Copyright (c) 2023-2025 EMQ Technologies Co., Ltd. All Rights Reserved.
+%%--------------------------------------------------------------------
+
+-module(emqx_cth_broker).
+
+-compile(export_all).
+-compile(nowarn_export_all).
+
+-spec connection_info(_Info, pid() | emqx_types:clientid()) -> _Value.
+connection_info(Info, Client) when is_pid(Client) ->
+    connection_info(Info, emqtt_info(clientid, Client));
+connection_info(Info, ClientId) ->
+    [ChanPid] = emqx_cm:lookup_channels(ClientId),
+    ConnMod = emqx_cm:do_get_chann_conn_mod(ClientId, ChanPid),
+    get_connection_info(Info, ConnMod, sys:get_state(ChanPid)).
+
+-spec connection_state(pid() | emqx_types:clientid()) -> _Value.
+connection_state(Client) when is_pid(Client) ->
+    connection_state(emqtt_info(clientid, Client));
+connection_state(ClientId) ->
+    [ChanPid] = emqx_cm:lookup_channels(ClientId),
+    ConnMod = emqx_cm:do_get_chann_conn_mod(ClientId, ChanPid),
+    ConnMod:get_state(ChanPid).
+
+get_connection_info(connmod, ConnMod, _State) ->
+    ConnMod;
+get_connection_info(Info, emqx_connection, State) ->
+    emqx_connection:info(Info, State);
+get_connection_info(Info, emqx_socket_connection, State) ->
+    emqx_socket_connection:info(Info, State);
+get_connection_info(Info, emqx_ws_connection, {_WSState, ConnState, _}) ->
+    emqx_ws_connection:info(Info, ConnState).
+
+emqtt_info(Key, Client) ->
+    proplists:get_value(Key, emqtt:info(Client), undefined).
@@ -181,9 +181,7 @@ t_tcp_frame_parsing_conn(_Config) ->
     with_listener(tcp, ?FUNCTION_NAME, Conf, fun() ->
         Client = emqtt_connect_tcp({127, 0, 0, 1}, Port),
         pong = emqtt:ping(Client),
-        ClientId = proplists:get_value(clientid, emqtt:info(Client)),
-        [CPid] = emqx_cm:lookup_channels(ClientId),
-        CState = emqx_connection:get_state(CPid),
+        CState = emqx_cth_broker:connection_state(Client),
         ?assertMatch(#{listener := {tcp, ?FUNCTION_NAME}}, CState),
         emqx_listeners:is_packet_parser_available(mqtt) andalso
             ?assertMatch(#{parser := {frame, _Options}}, CState)

@@ -235,7 +235,7 @@ test_change_parse_unit(ConfPath, ClientOpts) ->
     ?assertMatch({ok, _}, emqx:update_config(ConfPath, {update, ListenerRawConf1})),
     Client1 = emqtt_connect(ClientOpts),
     pong = emqtt:ping(Client1),
-    CState1 = get_conn_state(Client1),
+    CState1 = emqx_cth_broker:connection_state(Client1),
     emqx_listeners:is_packet_parser_available(mqtt) andalso
         ?assertMatch(
             #{parser := {frame, _Options}},
@@ -245,7 +245,7 @@ test_change_parse_unit(ConfPath, ClientOpts) ->
     ?assertMatch({ok, _}, emqx:update_config(ConfPath, {update, ListenerRawConf0})),
     Client2 = emqtt_connect(ClientOpts),
     pong = emqtt:ping(Client2),
-    CState2 = get_conn_state(Client2),
+    CState2 = emqx_cth_broker:connection_state(Client2),
     emqx_listeners:is_packet_parser_available(mqtt) andalso
         ?assertMatch(
             #{parser := Parser} when Parser =/= map_get(parser, CState1),
@@ -424,8 +424,3 @@ emqtt_connect(Opts) ->
         {error, Reason} ->
             error(Reason, [Opts])
     end.
-
-get_conn_state(Client) ->
-    ClientId = proplists:get_value(clientid, emqtt:info(Client)),
-    [CPid | _] = emqx_cm:lookup_channels(ClientId),
-    emqx_connection:get_state(CPid).
@@ -100,16 +100,6 @@ end_per_testcase(TestCase, Config) ->
 client_info(Key, Client) ->
     proplists:get_value(Key, emqtt:info(Client), undefined).
 
-connection_info(Info, ClientPid, Config) when is_list(Config) ->
-    connection_info(Info, ClientPid, ?config(conn_type, Config));
-connection_info(Info, ClientPid, tcp) ->
-    emqx_connection:info(Info, sys:get_state(ClientPid));
-connection_info(Info, ClientPid, quic) ->
-    emqx_connection:info(Info, sys:get_state(ClientPid));
-connection_info(Info, ClientPid, ws) ->
-    {_WSState, ConnState, _} = sys:get_state(ClientPid),
-    emqx_ws_connection:info(Info, ConnState).
-
 receive_messages(Count) ->
     receive_messages(Count, []).
 
@@ -283,8 +273,7 @@ t_connect_will_message(Config) ->
         | Config
     ]),
     {ok, _} = emqtt:ConnFun(Client1),
-    [ClientPid] = emqx_cm:lookup_channels(client_info(clientid, Client1)),
-    WillMsg = connection_info({channel, will_msg}, ClientPid, Config),
+    WillMsg = emqx_cth_broker:connection_info({channel, will_msg}, Client1),
     %% [MQTT-3.1.2-7]
     ?assertNotEqual(undefined, WillMsg),
 
@@ -468,13 +457,12 @@ t_connect_emit_stats_timeout(Config) ->
     {ok, _} = emqtt:ConnFun(Client),
     %% Poke the connection to ensure stats timer is armed.
     pong = emqtt:ping(Client),
-    [ClientPid] = emqx_cm:lookup_channels(client_info(clientid, Client)),
     ?assertMatch(
         TRef when is_reference(TRef),
-        connection_info(stats_timer, ClientPid, Config)
+        emqx_cth_broker:connection_info(stats_timer, Client)
     ),
     ?block_until(#{?snk_kind := cancel_stats_timer}, IdleTimeout * 2, _BackInTime = 0),
-    ?assertEqual(undefined, connection_info(stats_timer, ClientPid, Config)),
+    ?assertEqual(undefined, emqx_cth_broker:connection_info(stats_timer, Client)),
     ok = emqtt:disconnect(Client).
 
 %% [MQTT-3.1.2-22]

@@ -355,13 +355,12 @@ t_choose_impl(Config) ->
         | Config
     ]),
     {ok, _} = emqtt:ConnFun(Client),
-    [ChanPid] = emqx_cm:lookup_channels(ClientId),
     ?assertEqual(
         case ?config(persistence, Config) of
             false -> emqx_session_mem;
             ds -> emqx_persistent_session_ds
         end,
-        emqx_connection:info({channel, {session, impl}}, sys:get_state(ChanPid))
+        emqx_cth_broker:connection_info({channel, {session, impl}}, ClientId)
     ),
     ok = emqtt:disconnect(Client).
 

@@ -792,28 +792,28 @@ t_qos1_random_dispatch_if_all_members_are_down(Config) when is_list(Config) ->
     [Pid1, Pid2] = emqx_shared_sub:subscribers(Group, Topic),
     ?assert(is_process_alive(Pid1)),
     ?assert(is_process_alive(Pid2)),
-    ?retry(100, 10, ?assertEqual(disconnected, get_channel_info(conn_state, Pid1))),
-    ?retry(100, 10, ?assertEqual(disconnected, get_channel_info(conn_state, Pid2))),
+    ?retry(100, 10, ?assertEqual(disconnected, get_channel_info(conn_state, ClientId1))),
+    ?retry(100, 10, ?assertEqual(disconnected, get_channel_info(conn_state, ClientId2))),
 
     {ok, _} = emqtt:publish(ConnPub, Topic, <<"hello11">>, 1),
     ?retry(
         100,
         10,
         begin
-            Msgs1 = emqx_mqueue:to_list(get_mqueue(Pid1)),
-            Msgs2 = emqx_mqueue:to_list(get_mqueue(Pid2)),
+            Msgs1 = emqx_mqueue:to_list(get_mqueue(ClientId1)),
+            Msgs2 = emqx_mqueue:to_list(get_mqueue(ClientId2)),
             %% assert the message is in mqueue (because socket is closed)
             ?assertMatch([#message{payload = <<"hello11">>}], Msgs1 ++ Msgs2)
         end
     ),
     emqtt:stop(ConnPub),
     ok.
 
-get_mqueue(ConnPid) ->
-    get_channel_info({session, mqueue}, ConnPid).
+get_mqueue(Client) ->
+    get_channel_info({session, mqueue}, Client).
 
-get_channel_info(Info, ConnPid) ->
-    emqx_connection:info({channel, Info}, sys:get_state(ConnPid)).
+get_channel_info(Info, Client) ->
+    emqx_cth_broker:connection_info({channel, Info}, Client).
 
 %% No ack, QoS 2 subscriptions,
 %% client1 receives one message, send pubrec, then suspend

@@ -1090,10 +1090,13 @@ t_keepalive(Config) ->
     {ok, _} = emqtt:connect(C1),
     [Pid] = emqx_cm:lookup_channels(list_to_binary(ClientId)),
     %% will reset to max keepalive if keepalive > max keepalive
-    #{conninfo := #{keepalive := InitKeepalive}} = emqx_connection:info(Pid),
+    ?assertMatch(
+        #{conninfo := #{keepalive := InitKeepalive}},
+        emqx_cm:get_chan_info(list_to_binary(ClientId))
+    ),
     ?assertMatch(
         #{max_idle_millisecond := 65536500},
-        emqx_connection:info({channel, keepalive}, sys:get_state(Pid))
+        emqx_cth_broker:connection_info({channel, keepalive}, list_to_binary(ClientId))
     ),
 
     ?assertMatch(