Add test params for retained state registry

circuit-foundation/src/commonMain/kotlin/com/slack/circuit/foundation/internal/WithCompositionLocalProviders.kt

@@ -6,14 +6,16 @@ import androidx.compose.runtime.Composable
 import androidx.compose.runtime.InternalComposeApi
 import androidx.compose.runtime.ProvidedValue
 import androidx.compose.runtime.currentComposer
+import com.slack.circuit.runtime.InternalCircuitApi
 
 /**
  * A slightly more efficient version of [withCompositionLocalProvider] that only accepts a single
  * [value].
  */
 @Composable
+@InternalCircuitApi
 @OptIn(InternalComposeApi::class)
-internal fun <R> withCompositionLocalProvider(
+public fun <R> withCompositionLocalProvider(
   value: ProvidedValue<*>,
   content: @Composable () -> R,
 ): R {
@@ -30,8 +32,9 @@ internal fun <R> withCompositionLocalProvider(
  * @param content The content to provide the value to.
  */
 @Composable
+@InternalCircuitApi
 @OptIn(InternalComposeApi::class)
-internal fun <R> withCompositionLocalProvider(
+public fun <R> withCompositionLocalProvider(
   vararg values: ProvidedValue<*>,
   content: @Composable () -> R,
 ): R {

circuit-retained/src/commonMain/kotlin/com/slack/circuit/retained/RetainedStateRegistry.kt

@@ -96,7 +96,7 @@ internal class RetainedStateRegistryImpl(retained: MutableMap<String, List<Any?>
 
   override fun registerValue(key: String, valueProvider: RetainedValueProvider): Entry {
     require(key.isNotBlank()) { "Registered key is empty or blank" }
-    valueProviders.getOrPut(key) { mutableListOf() }.add(valueProvider)
+    valueProviders.getOrPut(key, ::mutableListOf).add(valueProvider)
     return object : Entry {
       override fun unregister() {
         val list = valueProviders.remove(key)

circuit-test/src/commonMain/kotlin/com/slack/circuit/test/PresenterTestExtensions.kt

@@ -9,7 +9,11 @@ import app.cash.molecule.RecompositionMode
 import app.cash.molecule.moleculeFlow
 import app.cash.turbine.ReceiveTurbine
 import app.cash.turbine.test
+import com.slack.circuit.foundation.internal.withCompositionLocalProvider
+import com.slack.circuit.retained.LocalRetainedStateRegistry
+import com.slack.circuit.retained.RetainedStateRegistry
 import com.slack.circuit.runtime.CircuitUiState
+import com.slack.circuit.runtime.InternalCircuitApi
 import com.slack.circuit.runtime.presenter.Presenter
 import kotlin.time.Duration
 
@@ -28,9 +32,10 @@ public suspend fun <UiState : CircuitUiState> Presenter<UiState>.test(
   timeout: Duration? = null,
   name: String? = null,
   policy: SnapshotMutationPolicy<UiState> = structuralEqualityPolicy(),
+  retainedStateRegistry: RetainedStateRegistry = RetainedStateRegistry(),
   block: suspend CircuitReceiveTurbine<UiState>.() -> Unit,
 ) {
-  presenterTestOf({ present() }, timeout, name, policy, block)
+  presenterTestOf({ present() }, timeout, name, policy, retainedStateRegistry, block)
 }
 
 /**
@@ -41,6 +46,7 @@ public suspend fun <UiState : CircuitUiState> Presenter<UiState>.test(
  * @param timeout an optional timeout for the test. Defaults to 1 second (in Turbine) if undefined.
  * @param policy a policy to controls how state changes are compared in
  *   [CircuitReceiveTurbine.awaitItem].
+ * @param retainedStateRegistry a [RetainedStateRegistry] that can operate
  * @param block the block to invoke.
  * @see moleculeFlow
  * @see test
@@ -50,9 +56,27 @@ public suspend fun <UiState : CircuitUiState> presenterTestOf(
   timeout: Duration? = null,
   name: String? = null,
   policy: SnapshotMutationPolicy<UiState> = structuralEqualityPolicy(),
+  retainedStateRegistry: RetainedStateRegistry = RetainedStateRegistry(),
   block: suspend CircuitReceiveTurbine<UiState>.() -> Unit,
 ) {
-  moleculeFlow(RecompositionMode.Immediate, presentFunction).test(timeout, name) {
-    asCircuitReceiveTurbine(policy).block()
+  try {
+    moleculeFlow(RecompositionMode.Immediate, decorate(presentFunction, retainedStateRegistry))
+      .test(timeout, name) { asCircuitReceiveTurbine(policy).block() }
+  } finally {
+    retainedStateRegistry.forgetUnclaimedValues()
   }
 }
+
+@OptIn(InternalCircuitApi::class)
+private fun <UiState : CircuitUiState> decorate(
+  presentFunction: @Composable () -> UiState,
+  retainedStateRegistry: RetainedStateRegistry,
+): @Composable () -> UiState {
+  val newFunction =
+    @Composable {
+      withCompositionLocalProvider(LocalRetainedStateRegistry provides retainedStateRegistry) {
+        presentFunction()
+      }
+    }
+  return newFunction
+}