Improve TestFormula to support key assertions.

formula-test/src/main/java/com/instacart/formula/test/TestFormula.kt

@@ -11,7 +11,7 @@ import java.lang.IllegalStateException
  * output updates and [inspect/interact][input] with input.
  */
 abstract class TestFormula<Input, Output> :
-    Formula<Input, TestFormula.State<Input, Output>, Output>() {
+    Formula<Input, TestFormula.State<Output>, Output>() {
 
     companion object {
         /**
@@ -29,9 +29,8 @@ abstract class TestFormula<Input, Output> :
         }
     }
 
-    data class State<Input, Output>(
-        val initialInput: Input,
-        val currentInput: Input,
+    data class State<Output>(
+        val key: Any?,
         val output: Output
     )
 
@@ -43,26 +42,18 @@ abstract class TestFormula<Input, Output> :
     /**
      * Uses initial input as key (to be decided if its robust enough)
      */
-    private val stateMap = mutableMapOf<Input, Value<Input, Output>>()
+    private val stateMap = mutableMapOf<Any?, Value<Input, Output>>()
 
     abstract fun initialOutput(): Output
 
-    override fun initialState(input: Input): State<Input, Output> {
-        return State(input, input, initialOutput())
+    override fun initialState(input: Input): State<Output> {
+        return State(key = key(input), output = initialOutput())
     }
 
-    override fun onInputChanged(
-        oldInput: Input,
-        input: Input,
-        state: State<Input, Output>
-    ): State<Input, Output> {
-        return state.copy(currentInput = input)
-    }
-
-    override fun Snapshot<Input, State<Input, Output>>.evaluate(): Evaluation<Output> {
-        stateMap[state.initialInput] = Value(
+    override fun Snapshot<Input, State<Output>>.evaluate(): Evaluation<Output> {
+        stateMap[state.key] = Value(
             input = input,
-            onNewOutput = context.onEvent<Output> {
+            onNewOutput = context.onEvent {
                 transition(state.copy(output = it))
             }
         )
@@ -72,7 +63,7 @@ abstract class TestFormula<Input, Output> :
             actions = context.actions {
                 Action.onTerminate().onEvent {
                     transition {
-                        stateMap.remove(state.initialInput)
+                        stateMap.remove(state.key)
                     }
                 }
             }
@@ -83,15 +74,36 @@ abstract class TestFormula<Input, Output> :
      * Emits a new [Output].
      */
     fun output(output: Output) {
-        val update = stateMap.values.lastOrNull()?.onNewOutput ?: throw IllegalStateException("Formula is not running")
+        val update = requireNotNull(stateMap.values.lastOrNull()?.onNewOutput) {
+            "Formula is not running"
+        }
         update(output)
     }
 
+    fun output(key: Any?, output: Output) {
+        val instance = requireNotNull(stateMap[key]) {
+            "Formula is not running"
+        }
+        instance.onNewOutput(output)
+    }
+
     /**
      * Performs an interaction on the current [Input] passed by the parent.
      */
     fun input(interact: Input.() -> Unit) {
-        val input = stateMap.values.lastOrNull()?.input ?: throw IllegalStateException("Formula is not running")
+        val input = requireNotNull(stateMap.values.lastOrNull()?.input) {
+            "Formula is not running"
+        }
         interact(input)
     }
+
+    /**
+     * Performs an interaction on the current [Input] passed by the parent.
+     */
+    fun input(key: Any?, interact: Input.() -> Unit) {
+        val instance = requireNotNull(stateMap[key]) {
+            "Formula for $key is not running, there are ${stateMap.keys} running"
+        }
+        instance.input.interact()
+    }
 }