Add LiveDataTestUtil

/*

 * Copyright (C) 2022 The Android Open Source Project

 *

 * Licensed under the Apache License, Version 2.0 (the "License");

 * you may not use this file except in compliance with the License.

 * You may obtain a copy of the License at

 *

 *      http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.android.settingslib.spa.testutils

import androidx.lifecycle.LiveData

import androidx.lifecycle.Observer

import java.util.concurrent.CountDownLatch

import java.util.concurrent.TimeUnit

import java.util.concurrent.TimeoutException

fun <T> LiveData<T>.getOrAwaitValue(

    timeout: Long = 1,

    timeUnit: TimeUnit = TimeUnit.SECONDS,

    afterObserve: () -> Unit = {},

): T? {

    var data: T? = null

    val latch = CountDownLatch(1)

    val observer = Observer<T> { newData ->

        data = newData

        latch.countDown()

    }

    this.observeForever(observer)

    afterObserve()

    try {

        // Don't wait indefinitely if the LiveData is not set.

        if (!latch.await(timeout, timeUnit)) {

            throw TimeoutException("LiveData value was never set.")

        }

    } finally {

        this.removeObserver(observer)

    }

    return data

}