Add modifyOutput function to TestFormula.

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

@@ -38,6 +38,7 @@ abstract class TestFormula<Input, Output> :
     data class Value<Input, Output>(
         val key: Any?,
         val input: Input,
+        val output: Output,
         val onNewOutput: (Output) -> Unit
     )
 
@@ -62,6 +63,7 @@ abstract class TestFormula<Input, Output> :
         stateMap[state.uniqueIdentifier] = Value(
             key = state.key,
             input = input,
+            output = state.output,
             onNewOutput = context.onEvent {
                 transition(state.copy(output = it))
             }
@@ -82,32 +84,40 @@ abstract class TestFormula<Input, Output> :
      * Emits a new [Output].
      */
     fun output(output: Output) {
-        val update = requireNotNull(stateMap.values.lastOrNull()?.onNewOutput) {
-            "Formula is not running"
-        }
+        val update = getMostRecentRunningFormula().onNewOutput
         update(output)
     }
 
     fun output(key: Any?, output: Output) {
-        val instance = getByKey(key)
+        val instance = getRunningFormulaByKey(key)
         instance.onNewOutput(output)
     }
 
+    fun updateOutput(modify: Output.() -> Output) {
+        val formulaValue = getMostRecentRunningFormula()
+        val newOutput = formulaValue.output.modify()
+        formulaValue.onNewOutput(newOutput)
+    }
+
+    fun updateOutput(key: Any?, modify: Output.() -> Output) {
+        val formulaValue = getRunningFormulaByKey(key)
+        val newOutput = formulaValue.output.modify()
+        formulaValue.onNewOutput(newOutput)
+    }
+
     /**
      * Performs an interaction on the current [Input] passed by the parent.
      */
     fun input(interact: Input.() -> Unit) {
-        val input = requireNotNull(stateMap.values.lastOrNull()?.input) {
-            "Formula is not running"
-        }
+        val input = getMostRecentRunningFormula().input
         interact(input)
     }
 
     /**
      * Performs an interaction on the current [Input] passed by the parent.
      */
     fun input(key: Any?, interact: Input.() -> Unit) {
-        val instance = getByKey(key)
+        val instance = getRunningFormulaByKey(key)
         instance.input.interact()
     }
 
@@ -118,7 +128,13 @@ abstract class TestFormula<Input, Output> :
         }
     }
 
-    private fun getByKey(key: Any?): Value<Input, Output> {
+    private fun getMostRecentRunningFormula(): Value<Input, Output> {
+        return requireNotNull(stateMap.values.lastOrNull()) {
+            "Formula is not running"
+        }
+    }
+
+    private fun getRunningFormulaByKey(key: Any?): Value<Input, Output> {
         return requireNotNull(stateMap.entries.firstOrNull { it.value.key == key }?.value) {
             val existingKeys = stateMap.entries.map { it.value.key }
             "Formula for $key is not running, there are $existingKeys running"

formula-test/src/test/java/com/instacart/formula/test/SimpleFormula.kt

@@ -0,0 +1,23 @@
+package com.instacart.formula.test
+
+import com.instacart.formula.IFormula
+import com.instacart.formula.test.SimpleFormula.Input
+import com.instacart.formula.test.SimpleFormula.Output
+
+interface SimpleFormula : IFormula<Input, Output> {
+    data class Input(val inputId: String = "inputId")
+    data class Output(val outputId: Int, val text: String)
+
+    override fun key(input: Input): Any? = "simple-formula-key"
+}
+
+class TestSimpleFormula(
+    private val initialOutput: Output = Output(
+        outputId = 0,
+        text = "",
+    )
+) : SimpleFormula {
+    override val implementation = testFormula(
+        initialOutput = initialOutput,
+    )
+}

formula-test/src/test/java/com/instacart/formula/test/TestFormulaTest.kt

@@ -6,108 +6,133 @@ import org.junit.Test
 
 class TestFormulaTest {
     @Test fun `assert running count is zero when formula is not running`() {
-        TestFormulaRobot()
-            .withTestFormula { assertRunningCount(0) }
-            .start()
-            .withTestFormula { assertRunningCount(1) }
+        val formula = TestSimpleFormula()
+        formula.implementation.assertRunningCount(0)
+        formula.test().input(SimpleFormula.Input())
+        formula.implementation.assertRunningCount(1)
+    }
+
+    @Test fun `emits initial output when subscribed`() {
+        val formula = TestSimpleFormula()
+        formula.test().input(SimpleFormula.Input()).output {
+            assertThat(this).isEqualTo(formula.implementation.initialOutput())
+        }
     }
 
     @Test fun `output throws an exception when no test formula is running`() {
         try {
-            TestFormulaRobot()
-                .withTestFormula { output(TestFormulaRobot.ChildFormula.Button(onNameChanged = {})) }
-
+            val formula = TestSimpleFormula()
+            formula.implementation.output(SimpleFormula.Output(0, ""))
             fail("Should not happen")
         } catch (e: Exception) {
             assertThat(e).hasMessageThat().startsWith("Formula is not running")
         }
     }
 
     @Test fun `output emits new output to the parent`() {
-        val newOutput = TestFormulaRobot.ChildFormula.Button(onNameChanged = {})
-        TestFormulaRobot()
-            .start()
-            .withTestFormula { output(newOutput) }
-            .assertOutput { assertThat(this.button).isEqualTo(newOutput) }
+        val formula = TestSimpleFormula()
+        val observer = formula.test().input(SimpleFormula.Input())
+
+        val newOutput = SimpleFormula.Output(5, "random")
+        formula.implementation.output(newOutput)
+        observer.output { assertThat(this).isEqualTo(newOutput) }
     }
 
     @Test fun `output with key throws an exception when test formula matching key is not running`() {
-        val newOutput = TestFormulaRobot.ChildFormula.Button(onNameChanged = {})
         try {
-            TestFormulaRobot()
-                .start()
-                .withTestFormula {
-                    output(key = "random-key", newOutput)
-                }
+            val newOutput = SimpleFormula.Output(5, "random")
 
+            val formula = TestSimpleFormula()
+            val observer = formula.test().input(SimpleFormula.Input())
+            formula.implementation.output("random-key", newOutput)
             fail("Should not happen")
         } catch (e: Exception) {
-            assertThat(e).hasMessageThat().startsWith("Formula for random-key is not running, there are [child-key] running")
+            assertThat(e).hasMessageThat().startsWith("Formula for random-key is not running, there are [simple-formula-key] running")
         }
     }
 
     @Test fun `output with key emits new output to the parent when key matches`() {
-        val newOutput = TestFormulaRobot.ChildFormula.Button(onNameChanged = {})
-        TestFormulaRobot()
-            .start()
-            .withTestFormula { output(key = "child-key", newOutput) }
-            .assertOutput { assertThat(this.button).isEqualTo(newOutput) }
+        val newOutput = SimpleFormula.Output(5, "random")
+
+        val formula = TestSimpleFormula()
+        val observer = formula.test().input(SimpleFormula.Input())
+        formula.implementation.output("simple-formula-key", newOutput)
+        observer.output { assertThat(this).isEqualTo(newOutput) }
+    }
+
+    @Test fun `updateOutput uses previous output and emits new one to the parent`() {
+        val formula = TestSimpleFormula()
+        val observer = formula.test().input(SimpleFormula.Input())
+
+        formula.implementation.updateOutput { copy(outputId = outputId.inc()) }
+        observer.output { assertThat(this).isEqualTo(SimpleFormula.Output(1, "")) }
+    }
+
+    @Test fun `updateOutput with key emits a modified output when key matches a running formula`() {
+        val formula = TestSimpleFormula()
+        val observer = formula.test().input(SimpleFormula.Input())
+
+        formula.implementation.updateOutput("simple-formula-key") { copy(outputId = outputId.inc()) }
+        observer.output { assertThat(this).isEqualTo(SimpleFormula.Output(1, "")) }
+    }
+
+    @Test fun `updateOutput with key throws an error when there is no running formula matching the key`() {
+        try {
+            val formula = TestSimpleFormula()
+            val observer = formula.test().input(SimpleFormula.Input())
+
+            formula.implementation.updateOutput("random-key") { copy(outputId = outputId.inc()) }
+            observer.output { assertThat(this).isEqualTo(SimpleFormula.Output(1, "")) }
+            fail()
+        } catch (e: Exception) {
+            assertThat(e).hasMessageThat().startsWith("Formula for random-key is not running, there are [simple-formula-key] running")
+        }
     }
 
     @Test fun `input() throw an exception when no formulas are running`() {
         try {
-            TestFormulaRobot()
-                .withTestFormula {
-                    input { onChangeName("my name") }
-                }
+            val formula = TestSimpleFormula()
+            formula.implementation.input {  }
 
             fail("Should not happen")
         } catch (e: Exception) {
             assertThat(e).hasMessageThat().startsWith("Formula is not running")
         }
     }
 
-    @Test fun `input() works as expected when formula is running`() {
-        TestFormulaRobot()
-            .start()
-            .withTestFormula {
-                input { onChangeName("my name") }
-            }
-            .assertOutput {
-                assertThat(name).isEqualTo("my name")
-            }
+    @Test fun `input() emits the last input provided by the parent`() {
+        val myInput = SimpleFormula.Input("my-input-id")
+        val formula = TestSimpleFormula()
+        val observer = formula.test()
+
+        // Initial input
+        observer.input(myInput)
+        formula.implementation.input { assertThat(this).isEqualTo(myInput) }
+
+        // Next input
+        val nextInput = SimpleFormula.Input("next-input-id")
+        observer.input(nextInput)
+        formula.implementation.input { assertThat(this).isEqualTo(nextInput) }
     }
 
     @Test fun `input() throws an error when NO formula that matches the key provided is running`() {
         try {
-            TestFormulaRobot()
-                .start()
-                .withTestFormula {
-                    input(key = "random-key") { onChangeName("my name") }
-                }
+            val formula = TestSimpleFormula()
+            formula.test().input(SimpleFormula.Input())
+            formula.implementation.input(key = "random-key") {}
 
             fail("Should not happen")
         } catch (e: Exception) {
-            assertThat(e).hasMessageThat().startsWith("Formula for random-key is not running, there are [child-key] running")
+            assertThat(e).hasMessageThat().startsWith("Formula for random-key is not running, there are [simple-formula-key] running")
         }
     }
 
     @Test fun `input() works as expected when formula that matches the key provided is running`() {
-        TestFormulaRobot()
-            .start()
-            .withTestFormula {
-                input(key = "child-key") { onChangeName("my name") }
-            }
-            .assertOutput {
-                assertThat(name).isEqualTo("my name")
-            }
-    }
-
-    @Test fun `input passed to formula`() {
-        TestFormulaRobot()
-            .start()
-            .withTestFormula {
-                input { assertThat(name).isEqualTo("") }
-            }
+        val myInput = SimpleFormula.Input()
+        val formula = TestSimpleFormula()
+        formula.test().input(myInput)
+        formula.implementation.input(key = "simple-formula-key") {
+            assertThat(this).isEqualTo(myInput)
+        }
     }
 }