https://issuetracker.google.com/issues/382294247

RealStateRestorationTester.kt

import android.annotation.SuppressLint
import android.os.Bundle
import android.os.Parcel
import android.os.Parcelable
import androidx.compose.runtime.Composable
import androidx.compose.runtime.CompositionLocalProvider
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.saveable.LocalSaveableStateRegistry
import androidx.compose.runtime.saveable.SaveableStateRegistry
import androidx.compose.runtime.setValue
import androidx.compose.ui.test.junit4.ComposeContentTestRule
import androidx.compose.ui.test.junit4.StateRestorationTester

/**
 * Like [StateRestorationTester], but actually performs state serialization to
 * catch errors that may occur on real devices but are missed by [StateRestorationTester]
 * due to [issue 382298310](https://issuetracker.google.com/issues/382298310) and
 * [issue 382294247](https://issuetracker.google.com/issues/382294247).
 */
class RealStateRestorationTester(private val composeTestRule: ComposeContentTestRule) {

  private var registry: RestorationRegistry? = null

  /** See [StateRestorationTester.setContent] */
  fun setContent(composable: @Composable () -> Unit) {
    composeTestRule.setContent {
      InjectRestorationRegistry { registry ->
        this.registry = registry
        composable()
      }
    }
  }

  /** See [StateRestorationTester.emulateSavedInstanceStateRestore] */
  fun emulateSavedInstanceStateRestore() {
    val registry = checkNotNull(registry) {
      "setContent should be called first!"
    }
    composeTestRule.runOnIdle {
      registry.saveStateAndDisposeChildren()
    }
    composeTestRule.runOnIdle {
      registry.emitChildrenWithRestoredState()
    }
    composeTestRule.runOnIdle {
      // we just wait for the children to be emitted
    }
  }

  @Composable
  @SuppressLint("ComposeUnstableReceiver")
  private fun InjectRestorationRegistry(content: @Composable (RestorationRegistry) -> Unit) {
    val original = requireNotNull(LocalSaveableStateRegistry.current) {
      "StateRestorationTester requires composeTestRule.setContent() to provide " +
        "a SaveableStateRegistry implementation via LocalSaveableStateRegistry"
    }
    val restorationRegistry = remember { RestorationRegistry(original) }
    CompositionLocalProvider(LocalSaveableStateRegistry provides restorationRegistry) {
      if (restorationRegistry.shouldEmitChildren) {
        content(restorationRegistry)
      }
    }
  }

  private class RestorationRegistry(private val original: SaveableStateRegistry) :
    SaveableStateRegistry {

    var shouldEmitChildren by mutableStateOf(true)
      private set
    private var currentRegistry: SaveableStateRegistry = original
    private lateinit var savedParcel: Parcel

    fun saveStateAndDisposeChildren() {
      savedParcel = Parcel.obtain().also {
        currentRegistry.performSave().toBundle().writeToParcel(it, 0)
      }
      shouldEmitChildren = false
    }

    fun emitChildrenWithRestoredState() {
      currentRegistry = SaveableStateRegistry(
        restoredValues = run {
          savedParcel.setDataPosition(0)
          Bundle.CREATOR.createFromParcel(savedParcel).toMap()
        },
        canBeSaved = { original.canBeSaved(it) }
      )
      shouldEmitChildren = true
    }

    override fun consumeRestored(key: String) = currentRegistry.consumeRestored(key)

    override fun registerProvider(key: String, valueProvider: () -> Any?) =
      currentRegistry.registerProvider(key, valueProvider)

    override fun canBeSaved(value: Any) = currentRegistry.canBeSaved(value)

    override fun performSave() = currentRegistry.performSave()
  }
}

// Copied from DisposableSaveableStateRegistry.android.kt
@Suppress("DEPRECATION", "UNCHECKED_CAST")
private fun Bundle.toMap(): Map<String, List<Any?>>? {
  val map = mutableMapOf<String, List<Any?>>()
  this.keySet().forEach { key ->
    val list = getParcelableArrayList<Parcelable?>(key) as ArrayList<Any?>
    map[key] = list
  }
  return map
}

// Copied from DisposableSaveableStateRegistry.android.kt
@Suppress("UNCHECKED_CAST")
private fun Map<String, List<Any?>>.toBundle(): Bundle {
  val bundle = Bundle()
  forEach { (key, list) ->
    val arrayList = if (list is ArrayList<Any?>) list else ArrayList(list)
    bundle.putParcelableArrayList(
      key,
      arrayList as ArrayList<Parcelable?>
    )
  }
  return bundle
}