test: refactor tests to use shared structure

apps/emqx_bridge_azure_blob_storage/test/emqx_bridge_azure_blob_storage_SUITE.erl

@@ -15,6 +15,7 @@
 -include_lib("snabbkaffe/include/snabbkaffe.hrl").
 -include_lib("erlazure/include/erlazure.hrl").
 -include("../src/emqx_bridge_azure_blob_storage.hrl").
+-include_lib("emqx_utils/include/emqx_message.hrl").
 
 -define(ACCOUNT_NAME_BIN, <<"devstoreaccount1">>).
 -define(ACCOUNT_KEY_BIN, <<
@@ -523,78 +524,49 @@ t_aggreg_upload_restart(Config) ->
 %% so while preserving uncompromised data.
 t_aggreg_upload_restart_corrupted(Config) ->
     ActionName = ?config(action_name, Config),
-    AggregId = aggreg_id(ActionName),
     BatchSize = ?CONF_MAX_RECORDS div 2,
-    ?check_trace(
-        #{timetrap => timer:seconds(30)},
-        begin
-            %% Create a bridge with the sample configuration.
-            ?assertMatch({ok, _Bridge}, emqx_bridge_v2_testlib:create_bridge_api(Config)),
-            {ok, _Rule} =
-                emqx_bridge_v2_testlib:create_rule_and_action_http(
-                    ?ACTION_TYPE_BIN, <<"">>, Config, #{
-                        sql => <<
-                            "SELECT"
-                            "  *,"
-                            "  strlen(payload) as psize,"
-                            "  unix_ts_to_rfc3339(publish_received_at, 'millisecond') as publish_received_at"
-                            "  FROM 'abs/#'"
-                        >>
-                    }
-                ),
-            Messages1 = [
-                {integer_to_binary(N), <<"abs/a/b/c">>, <<"{\"hello\":\"world\"}">>}
-             || N <- lists:seq(1, BatchSize)
-            ],
-            %% Ensure that they span multiple batch queries.
-            {ok, {ok, _}} =
-                ?wait_async_action(
-                    publish_messages_delayed(lists:map(fun mk_message/1, Messages1), 1),
-                    #{?snk_kind := connector_aggreg_records_written, action := AggregId}
-                ),
-            ct:pal("first batch's records have been written"),
-
-            %% Find out the buffer file.
-            {ok, #{filename := Filename}} = ?block_until(
-                #{?snk_kind := connector_aggreg_buffer_allocated, action := AggregId}
-            ),
-            ct:pal("new buffer allocated"),
-
-            %% Stop the bridge, corrupt the buffer file, and restart the bridge.
-            {ok, {{_, 204, _}, _, _}} = emqx_bridge_v2_testlib:disable_kind_http_api(Config),
-            BufferFileSize = filelib:file_size(Filename),
-            ok = emqx_connector_aggregator_test_helpers:truncate_at(Filename, BufferFileSize div 2),
-            {ok, {{_, 204, _}, _, _}} = emqx_bridge_v2_testlib:enable_kind_http_api(Config),
-
-            %% Send some more messages.
-            Messages2 = [
+    Opts = #{
+        aggreg_id => aggreg_id(ActionName),
+        batch_size => BatchSize,
+        rule_sql => <<
+            "SELECT"
+            "  *,"
+            "  strlen(payload) as psize,"
+            "  unix_ts_to_rfc3339(publish_received_at, 'millisecond') as publish_received_at"
+            "  FROM 'abs/#'"
+        >>,
+        make_message_fn => fun(N) ->
+            mk_message(
                 {integer_to_binary(N), <<"abs/a/b/c">>, <<"{\"hello\":\"world\"}">>}
-             || N <- lists:seq(1, BatchSize)
-            ],
-            ok = publish_messages_delayed(lists:map(fun mk_message/1, Messages2), 1),
-            ct:pal("published second batch"),
-
-            %% Wait until the delivery is completed.
-            {ok, _} = ?block_until(#{
-                ?snk_kind := connector_aggreg_delivery_completed, action := AggregId
-            }),
-            ct:pal("delivery completed"),
-
-            %% Check that upload contains part of the first batch and all of the second batch.
+            )
+        end,
+        message_check_fn => fun(Context) ->
+            #{
+                messages_before := Messages1,
+                messages_after := Messages2
+            } = Context,
             [#cloud_blob{name = BlobName}] = list_blobs(Config),
             {ok, CSV = [_Header | Rows]} = erl_csv:decode(get_blob(BlobName, Config)),
             NRows = length(Rows),
             ?assert(NRows > BatchSize, CSV),
+            Expected = [
+                {ClientId, Topic, Payload}
+             || #message{
+                    from = ClientId,
+                    topic = Topic,
+                    payload = Payload
+                } <- lists:sublist(Messages1, NRows - BatchSize) ++ Messages2
+            ],
             ?assertEqual(
-                lists:sublist(Messages1, NRows - BatchSize) ++ Messages2,
+                Expected,
                 [{ClientID, Topic, Payload} || [_TS, ClientID, Topic, Payload | _] <- Rows],
                 CSV
             ),
 
             ok
-        end,
-        []
-    ),
+        end
+    },
+    emqx_bridge_v2_testlib:t_aggreg_upload_restart_corrupted(Config, Opts),
     ok.
 
 %% This test verifies that the bridge will finish uploading a buffer file after a restart.

apps/emqx_bridge_couchbase/test/emqx_bridge_couchbase_SUITE.erl

@@ -366,11 +366,6 @@ get_all_rows(Scope, Collection, Config) ->
             Resp
     end.
 
-proplist_update(Proplist, K, Fn) ->
-    {K, OldV} = lists:keyfind(K, 1, Proplist),
-    NewV = Fn(OldV),
-    lists:keystore(K, 1, Proplist, {K, NewV}).
-
 pre_publish_fn(Scope, Collection, Config) ->
     fun(Context) ->
         ensure_scope(Scope, Config),
@@ -417,7 +412,7 @@ t_rule_action(Config) ->
 
 %% batch is not yet supported
 skip_t_rule_action_batch(Config0) ->
-    Config = proplist_update(Config0, action_config, fun(ActionConfig) ->
+    Config = emqx_bridge_v2_testlib:proplist_update(Config0, action_config, fun(ActionConfig) ->
         emqx_utils_maps:deep_merge(
             ActionConfig,
             #{

apps/emqx_bridge_s3/test/emqx_bridge_s3_aggreg_upload_SUITE.erl

@@ -10,6 +10,7 @@
 -include_lib("eunit/include/eunit.hrl").
 -include_lib("common_test/include/ct.hrl").
 -include_lib("snabbkaffe/include/test_macros.hrl").
+-include_lib("emqx_utils/include/emqx_message.hrl").
 
 -import(emqx_utils_conv, [bin/1]).
 
@@ -351,58 +352,60 @@ t_aggreg_upload_restart(Config) ->
         erl_csv:decode(Content)
     ).
 
+%% NOTE
+%% This test verifies that the bridge can recover from a buffer file corruption,
+%% and does so while preserving uncompromised data.
 t_aggreg_upload_restart_corrupted(Config) ->
-    %% NOTE
-    %% This test verifies that the bridge can recover from a buffer file corruption,
-    %% and does so while preserving uncompromised data.
     Bucket = ?config(s3_bucket, Config),
     BridgeName = ?config(bridge_name, Config),
-    AggregId = aggreg_id(BridgeName),
     BatchSize = ?CONF_MAX_RECORDS div 2,
-    %% Create a bridge with the sample configuration.
-    ?assertMatch({ok, _Bridge}, emqx_bridge_v2_testlib:create_bridge(Config)),
-    %% Send some sample messages that look like Rule SQL productions.
-    Messages1 = [
-        {integer_to_binary(N), <<"a/b/c">>, <<"{\"hello\":\"world\"}">>}
-     || N <- lists:seq(1, BatchSize)
-    ],
-    %% Ensure that they span multiple batch queries.
-    ok = send_messages_delayed(BridgeName, lists:map(fun mk_message_event/1, Messages1), 1),
-    {ok, _} = ?block_until(
-        #{?snk_kind := connector_aggreg_records_written, action := AggregId},
-        infinity,
-        0
-    ),
-    %% Find out the buffer file.
-    {ok, #{filename := Filename}} = ?block_until(
-        #{?snk_kind := connector_aggreg_buffer_allocated, action := AggregId}
-    ),
-    %% Stop the bridge, corrupt the buffer file, and restart the bridge.
-    {ok, _} = emqx_bridge_v2:disable_enable(disable, ?BRIDGE_TYPE, BridgeName),
-    BufferFileSize = filelib:file_size(Filename),
-    ok = emqx_connector_aggregator_test_helpers:truncate_at(Filename, BufferFileSize div 2),
-    {ok, _} = emqx_bridge_v2:disable_enable(enable, ?BRIDGE_TYPE, BridgeName),
-    %% Send some more messages.
-    Messages2 = [
-        {integer_to_binary(N), <<"c/d/e">>, <<"{\"hello\":\"world\"}">>}
-     || N <- lists:seq(1, BatchSize)
-    ],
-    ok = send_messages_delayed(BridgeName, lists:map(fun mk_message_event/1, Messages2), 0),
-    %% Wait until the delivery is completed.
-    {ok, _} = ?block_until(#{?snk_kind := connector_aggreg_delivery_completed, action := AggregId}),
-    %% Check that upload contains part of the first batch and all of the second batch.
-    _Uploads = [#{key := Key}] = emqx_bridge_s3_test_helpers:list_objects(Bucket),
-    CSV = [_Header | Rows] = fetch_parse_csv(Bucket, Key),
-    NRows = length(Rows),
-    ?assert(
-        NRows > BatchSize,
-        CSV
-    ),
-    ?assertEqual(
-        lists:sublist(Messages1, NRows - BatchSize) ++ Messages2,
-        [{ClientID, Topic, Payload} || [_TS, ClientID, Topic, Payload | _] <- Rows],
-        CSV
-    ).
+    Opts = #{
+        aggreg_id => aggreg_id(BridgeName),
+        batch_size => BatchSize,
+        rule_sql => <<
+            "SELECT"
+            "  *,"
+            "  strlen(payload) as psize,"
+            "  unix_ts_to_rfc3339(publish_received_at, 'millisecond') as publish_received_at"
+            "  FROM 's3/#'"
+        >>,
+        make_message_fn => fun(N) ->
+            mk_message(
+                {integer_to_binary(N), <<"s3/a/b/c">>, <<"{\"hello\":\"world\"}">>}
+            )
+        end,
+        message_check_fn => fun(Context) ->
+            #{
+                messages_before := Messages1,
+                messages_after := Messages2
+            } = Context,
+
+            _Uploads = [#{key := Key}] = emqx_bridge_s3_test_helpers:list_objects(Bucket),
+            CSV = [_Header | Rows] = fetch_parse_csv(Bucket, Key),
+            NRows = length(Rows),
+            ?assert(
+                NRows > BatchSize,
+                CSV
+            ),
+            Expected = [
+                {ClientId, Topic, Payload}
+             || #message{
+                    from = ClientId,
+                    topic = Topic,
+                    payload = Payload
+                } <- lists:sublist(Messages1, NRows - BatchSize) ++ Messages2
+            ],
+            ?assertEqual(
+                Expected,
+                [{ClientID, Topic, Payload} || [_TS, ClientID, Topic, Payload | _] <- Rows],
+                CSV
+            ),
+
+            ok
+        end
+    },
+    emqx_bridge_v2_testlib:t_aggreg_upload_restart_corrupted(Config, Opts),
+    ok.
 
 t_aggreg_pending_upload_restart(Config) ->
     %% NOTE
@@ -511,6 +514,9 @@ receive_sender_reports([]) ->
 
 %%
 
+mk_message({ClientId, Topic, Payload}) ->
+    emqx_message:make(bin(ClientId), bin(Topic), Payload).
+
 mk_message_event({ClientID, Topic, Payload}) ->
     emqx_bridge_s3_test_helpers:mk_message_event(ClientID, Topic, Payload).