diff --git a/big_tests/dynamic_domains.spec b/big_tests/dynamic_domains.spec
index 90fdd1737b..6535923117 100644
--- a/big_tests/dynamic_domains.spec
+++ b/big_tests/dynamic_domains.spec
@@ -80,6 +80,9 @@
 
 {suites, "tests", vcard_SUITE}.
 {suites, "tests", vcard_simple_SUITE}.
+
+{suites, "tests", xep_0352_csi_SUITE}.
+
 {suites, "tests", domain_removal_SUITE}.
 
 {config, ["dynamic_domains.config", "test.config"]}.
diff --git a/big_tests/tests/xep_0352_csi_SUITE.erl b/big_tests/tests/xep_0352_csi_SUITE.erl
index b80f5a8e82..922fabbf4a 100644
--- a/big_tests/tests/xep_0352_csi_SUITE.erl
+++ b/big_tests/tests/xep_0352_csi_SUITE.erl
@@ -5,6 +5,8 @@
 
 -compile([export_all]).
 
+-import(domain_helper, [host_type/0]).
+
 -define(CSI_BUFFER_MAX, 10).
 
 all() ->
@@ -12,8 +14,7 @@ all() ->
 
 
 groups() ->
-    G = [{basic, [parallel, shuffle], all_tests()}],
-    ct_helper:repeat_all_until_all_ok(G).
+    [{basic, [parallel, shuffle], all_tests()}].
 
 all_tests() ->
     [
@@ -32,13 +33,11 @@ suite() ->
     escalus:suite().
 
 init_per_suite(Config) ->
-    Domain = ct:get_config({hosts, mim, domain}),
-    dynamic_modules:start(Domain, mod_csi, [{buffer_max, ?CSI_BUFFER_MAX}]),
+    dynamic_modules:start(host_type(), mod_csi, [{buffer_max, ?CSI_BUFFER_MAX}]),
     [{escalus_user_db, {module, escalus_ejabberd}} | escalus:init_per_suite(Config)].
 
 end_per_suite(Config) ->
-    Domain = ct:get_config({hosts, mim, domain}),
-    dynamic_modules:stop(Domain, mod_csi),
+    dynamic_modules:stop(host_type(), mod_csi),
     escalus_fresh:clean(),
     escalus:end_per_suite(Config).
 
diff --git a/src/mod_csi.erl b/src/mod_csi.erl
index 93278570d2..080ae509df 100644
--- a/src/mod_csi.erl
+++ b/src/mod_csi.erl
@@ -7,9 +7,13 @@
 -behaviour(gen_mod).
 -behaviour(mongoose_module_metrics).
 
--export([start/2]).
--export([stop/1]).
--export([config_spec/0]).
+%% gen_mod callbacks
+-export([start/2, 
+         stop/1, 
+         config_spec/0, 
+         supported_features/0]).
+
+%% Hook handlers 
 -export([c2s_stream_features/3]).
 
 -ignore_xref([c2s_stream_features/3]).
@@ -21,17 +25,28 @@
 
 -export_type([state/0]).
 
-start(Host, _Opts) ->
-    [ejabberd_hooks:add(Name, Host, Module, Function, Priority) ||
-     {Name, Module, Function, Priority} <- hooks()],
-    ensure_metrics(Host),
+-spec start(mongooseim:host_type(), gen_mod:module_opts()) -> ok.
+start(HostType, _Opts) ->
+    ejabberd_hooks:add(hooks(HostType)),
+    ensure_metrics(HostType),
     ok.
 
-stop(Host) ->
-    [ejabberd_hooks:delete(Name, Host, Module, Function, Priority) ||
-     {Name, Module, Function, Priority} <- hooks()],
+-spec stop(mongooseim:host_type()) -> ok.
+stop(HostType) ->
+    ejabberd_hooks:delete(hooks(HostType)),
     ok.
 
+hooks(HostType) ->
+    [{c2s_stream_features, HostType, ?MODULE, c2s_stream_features, 60}].
+
+ensure_metrics(HostType) ->
+    mongoose_metrics:ensure_metric(HostType, [HostType, modCSIInactive], spiral),
+    mongoose_metrics:ensure_metric(HostType, [HostType, modCSIInactive], spiral).
+
+%%%
+%%% config_spec
+%%%
+
 -spec config_spec() -> mongoose_config_spec:config_section().
 config_spec() ->
     #section{
@@ -39,12 +54,13 @@ config_spec() ->
                                              validate = non_negative}}
       }.
 
-hooks() ->
-    [{c2s_stream_features, ?MODULE, c2s_stream_features, 60}].
+-spec supported_features() -> [atom()].
+supported_features() ->
+    [dynamic_domains].
 
-ensure_metrics(Host) ->
-    mongoose_metrics:ensure_metric(Host, [Host, modCSIInactive], spiral),
-    mongoose_metrics:ensure_metric(Host, [Host, modCSIInactive], spiral).
+%%%
+%%% Hook handlers
+%%%
 
 -spec c2s_stream_features([exml:element()], mongooseim:host_type(), jid:lserver()) ->
           [exml:element()].