Extend function to do unit test with Firebase realtime database

RealTimeDatabaseTest.kt

import androidx.test.ext.junit.runners.AndroidJUnit4
import com.google.android.gms.tasks.Task
import com.google.common.truth.Truth.assertThat
import com.google.firebase.database.DataSnapshot
import com.google.firebase.database.DatabaseError
import com.google.firebase.database.DatabaseReference
import com.google.firebase.database.ValueEventListener
import com.google.firebase.database.ktx.database
import com.google.firebase.ktx.Firebase
import org.junit.Before
import org.junit.Test
import org.junit.runner.RunWith
import java.util.concurrent.CountDownLatch


fun DatabaseReference.blockSetValue(value: Any): Task<Void> = setValue(value).apply {
    val completedSignal = CountDownLatch(1)
    addOnCompleteListener {
        completedSignal.countDown()
    }
    completedSignal.await()
}


fun DatabaseReference.addListenerForSingleValueEvent(
    valueEventListener: ValueEventListener,
    wait: Boolean = false
) = if (wait) {
    addListenerForSingleValueEvent(valueEventListener)
    val changedSignal = CountDownLatch(1)
    addListenerForSingleValueEvent(object : ValueEventListener {
        override fun onDataChange(snapshot: DataSnapshot) {
            changedSignal.countDown()
        }

        override fun onCancelled(error: DatabaseError) {
            valueEventListener.onCancelled(error)
        }
    })
    changedSignal.await()
} else {
    addListenerForSingleValueEvent(valueEventListener)
}


@RunWith(AndroidJUnit4::class)
class RealTimeDatabase {

    companion object {
        const val messageContent = "Hello world!"
        const val messagePath = "test"
    }

    private lateinit var messageRef: DatabaseReference

    @Before
    fun setup() {
        messageRef = Firebase.database.getReference(messagePath)
    }

    @Test
    fun writeToDatabase() {
        messageRef.blockSetValue(messageContent)
    }

    @Test
    fun readFromDatabase() {
        messageRef
            .addListenerForSingleValueEvent(object : ValueEventListener {
                override fun onDataChange(snapshot: DataSnapshot) {
                    val value = snapshot.value
                    assertThat(value).isEqualTo(messageContent)
                }

                override fun onCancelled(error: DatabaseError) {
                }
            }, true)
    }
}

Add waiting for write/read realtime database finished when do unit test.