diff --git a/stream-chat-android-client/src/main/java/io/getstream/chat/android/client/ChatClient.kt b/stream-chat-android-client/src/main/java/io/getstream/chat/android/client/ChatClient.kt
index 59e7f07936f..380e6d293b8 100644
--- a/stream-chat-android-client/src/main/java/io/getstream/chat/android/client/ChatClient.kt
+++ b/stream-chat-android-client/src/main/java/io/getstream/chat/android/client/ChatClient.kt
@@ -372,12 +372,11 @@ internal constructor(
             )
     }
 
-    @PublishedApi
     @InternalStreamChatApi
     @StreamHandsOff(
         "This method is used to avoid race-condition between plugin initialization and dependency resolution.",
     )
-    internal fun awaitInitializationState(timeoutMilliseconds: Long): InitializationState? {
+    public fun awaitInitializationState(timeoutMilliseconds: Long): InitializationState? {
         var initState: InitializationState? = clientState.initializationState.value
         var spendTime = 0L
         inheritScope { Job(it) }.launch {
diff --git a/stream-chat-android-compose/src/test/kotlin/io/getstream/chat/android/compose/viewmodel/channels/ChannelListViewModelTest.kt b/stream-chat-android-compose/src/test/kotlin/io/getstream/chat/android/compose/viewmodel/channels/ChannelListViewModelTest.kt
index 1c41b433ba8..4526b74d9d8 100644
--- a/stream-chat-android-compose/src/test/kotlin/io/getstream/chat/android/compose/viewmodel/channels/ChannelListViewModelTest.kt
+++ b/stream-chat-android-compose/src/test/kotlin/io/getstream/chat/android/compose/viewmodel/channels/ChannelListViewModelTest.kt
@@ -330,6 +330,7 @@ internal class ChannelListViewModelTest {
         fun givenCurrentUser(currentUser: User = User(id = "Jc")) = apply {
             whenever(clientState.user) doReturn MutableStateFlow(currentUser)
             whenever(clientState.initializationState) doReturn MutableStateFlow(InitializationState.COMPLETE)
+            whenever(chatClient.awaitInitializationState(any())) doReturn InitializationState.COMPLETE
         }
 
         fun givenChannelMutes(channelMutes: List<ChannelMute> = emptyList()) = apply {
diff --git a/stream-chat-android-compose/src/test/kotlin/io/getstream/chat/android/compose/viewmodel/messages/MessageComposerViewModelTest.kt b/stream-chat-android-compose/src/test/kotlin/io/getstream/chat/android/compose/viewmodel/messages/MessageComposerViewModelTest.kt
index f6c0b588fe7..fb21627bfce 100644
--- a/stream-chat-android-compose/src/test/kotlin/io/getstream/chat/android/compose/viewmodel/messages/MessageComposerViewModelTest.kt
+++ b/stream-chat-android-compose/src/test/kotlin/io/getstream/chat/android/compose/viewmodel/messages/MessageComposerViewModelTest.kt
@@ -402,6 +402,7 @@ internal class MessageComposerViewModelTest {
             whenever(clientState.user) doReturn MutableStateFlow(currentUser)
             whenever(chatClient.clientState) doReturn clientState
             whenever(clientState.initializationState) doReturn MutableStateFlow(InitializationState.COMPLETE)
+            whenever(chatClient.awaitInitializationState(any())) doReturn InitializationState.COMPLETE
         }
 
         fun givenChannelQuery(channel: Channel = Channel()) = apply {
diff --git a/stream-chat-android-state/src/test/java/io/getstream/chat/android/state/extensions/ChatClientExtensionTests.kt b/stream-chat-android-state/src/test/java/io/getstream/chat/android/state/extensions/ChatClientExtensionTests.kt
index 10a14961375..13fa76f8cd5 100644
--- a/stream-chat-android-state/src/test/java/io/getstream/chat/android/state/extensions/ChatClientExtensionTests.kt
+++ b/stream-chat-android-state/src/test/java/io/getstream/chat/android/state/extensions/ChatClientExtensionTests.kt
@@ -46,6 +46,7 @@ import org.amshove.kluent.shouldBeEqualTo
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.extension.RegisterExtension
+import org.mockito.kotlin.any
 import org.mockito.kotlin.doAnswer
 import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.mock
@@ -136,6 +137,7 @@ internal class ChatClientExtensionTests {
                 userFlow.value = null
                 Unit.asCall()
             }
+            on(it.awaitInitializationState(any())) doReturn InitializationState.COMPLETE
         }
     }
 
diff --git a/stream-chat-android-ui-components/src/test/kotlin/io/getstream/chat/android/ui/viewmodels/channels/ChannelListViewModelTest.kt b/stream-chat-android-ui-components/src/test/kotlin/io/getstream/chat/android/ui/viewmodels/channels/ChannelListViewModelTest.kt
index d6572ad02e4..1d74186b865 100644
--- a/stream-chat-android-ui-components/src/test/kotlin/io/getstream/chat/android/ui/viewmodels/channels/ChannelListViewModelTest.kt
+++ b/stream-chat-android-ui-components/src/test/kotlin/io/getstream/chat/android/ui/viewmodels/channels/ChannelListViewModelTest.kt
@@ -264,6 +264,7 @@ internal class ChannelListViewModelTest {
             whenever(clientState.user) doReturn MutableStateFlow(currentUser)
             whenever(clientState.initializationState) doReturn MutableStateFlow(InitializationState.COMPLETE)
             whenever(chatClient.getCurrentUser()) doReturn currentUser
+            whenever(chatClient.awaitInitializationState(any())) doReturn InitializationState.COMPLETE
         }
 
         fun givenChannelMutes(channelMutes: List<ChannelMute> = emptyList()) = apply {
diff --git a/stream-chat-android-ui-components/src/test/kotlin/io/getstream/chat/android/ui/viewmodels/messages/MessageComposerViewModelTest.kt b/stream-chat-android-ui-components/src/test/kotlin/io/getstream/chat/android/ui/viewmodels/messages/MessageComposerViewModelTest.kt
index c4b6a22d595..2145b206838 100644
--- a/stream-chat-android-ui-components/src/test/kotlin/io/getstream/chat/android/ui/viewmodels/messages/MessageComposerViewModelTest.kt
+++ b/stream-chat-android-ui-components/src/test/kotlin/io/getstream/chat/android/ui/viewmodels/messages/MessageComposerViewModelTest.kt
@@ -403,6 +403,7 @@ internal class MessageComposerViewModelTest {
             whenever(clientState.user) doReturn MutableStateFlow(currentUser)
             whenever(chatClient.clientState) doReturn clientState
             whenever(clientState.initializationState) doReturn MutableStateFlow(InitializationState.COMPLETE)
+            whenever(chatClient.awaitInitializationState(any())) doReturn InitializationState.COMPLETE
         }
 
         fun givenChannelQuery(channel: Channel = Channel()) = apply {