diff --git a/packages/SettingsLib/Spa/testutils/src/com/android/settingslib/spa/testutils/LiveDataTestUtil.kt b/packages/SettingsLib/Spa/testutils/src/com/android/settingslib/spa/testutils/LiveDataTestUtil.kt new file mode 100644 index 0000000000000..dddda5511c7bd --- /dev/null +++ b/packages/SettingsLib/Spa/testutils/src/com/android/settingslib/spa/testutils/LiveDataTestUtil.kt @@ -0,0 +1,50 @@ +/* + * 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 LiveData.getOrAwaitValue( + timeout: Long = 1, + timeUnit: TimeUnit = TimeUnit.SECONDS, + afterObserve: () -> Unit = {}, +): T? { + var data: T? = null + val latch = CountDownLatch(1) + val observer = Observer { 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 +}