diff --git a/formula-test/src/main/java/com/instacart/formula/test/TestExtensions.kt b/formula-test/src/main/java/com/instacart/formula/test/TestExtensions.kt
index 61c993cb..3b630eec 100644
--- a/formula-test/src/main/java/com/instacart/formula/test/TestExtensions.kt
+++ b/formula-test/src/main/java/com/instacart/formula/test/TestExtensions.kt
@@ -40,8 +40,22 @@ fun <Input : Any, State : Any, Output : Any> withSnapshot(
     return observer.values().last()
 }
 
+/**
+ *
+ * Creates a test formula for a specific [IFormula] instance.
+ *
+ * ```kotlin
+ * class FakeUserFormula : UserFormula {
+ *      override val implementation = testFormula(
+ *          initialOutput = User(
+ *              name = "Fake user name",
+ *          )
+ *      )
+ * }
+ * ```
+ */
 fun <Input, Output> IFormula<Input, Output>.testFormula(
     initialOutput: Output,
 ): TestFormula<Input, Output> {
-    return TestFormula(initialOutput, key = { this.key(it) })
+    return TestFormula(this, initialOutput)
 }
\ No newline at end of file
diff --git a/formula-test/src/main/java/com/instacart/formula/test/TestFormula.kt b/formula-test/src/main/java/com/instacart/formula/test/TestFormula.kt
index a517441c..ca0d561c 100644
--- a/formula-test/src/main/java/com/instacart/formula/test/TestFormula.kt
+++ b/formula-test/src/main/java/com/instacart/formula/test/TestFormula.kt
@@ -3,69 +3,73 @@ package com.instacart.formula.test
 import com.instacart.formula.Action
 import com.instacart.formula.Evaluation
 import com.instacart.formula.Formula
+import com.instacart.formula.IFormula
 import com.instacart.formula.Snapshot
+import com.instacart.formula.test.utils.RunningInstanceManager
 import java.util.concurrent.atomic.AtomicLong
 
 /**
- * Test formula is used to provide a fake formula implementation. It allows you to [send][output]
- * output updates and [inspect/interact][input] with input.
+ * Test formula is used to replace a real formula with a fake implementation. This allows you to:
+ * - Verify that parent passes correct inputs. Take a look at [input].
+ * - Verify that parent deals with output changes correctly. Take a look at [output]
+ *
+ *
+ * ```kotlin
+ * // To replace a real formula with a fake one, we need first define the interface
+ * interface UserFormula : IFormula<Input, Output> {
+ *     data class Input(val userId: String)
+ *     data class Output(val user: User?)
+ * }
+ *
+ * // Then, we can create a fake implementation
+ * class FakeUserFormula : UserFormula {
+ *     override val implementation = testFormula(
+ *         initialOutput = Output(user = null)
+ *     )
+ * }
+ *
+ * // Then, in our test, we can do the following
+ * @Test fun `ensure that account formula passes user id to user formula`() {
+ *     val userFormula = FakeUserFormula()
+ *     val accountFormula = RealAccountFormula(userFormula)
+ *     accountFormula.test().input(AccountFormula.Input(userId = "my-user-id"))
+ *
+ *     userFormula.implementation.input {
+ *       Truth.assertThat(this.userId).isEqualTo("my-user-id")
+ *     }
+ * }
+ * ```
  */
-abstract class TestFormula<Input, Output> :
-    Formula<Input, TestFormula.State<Output>, Output>() {
-
-    companion object {
-        /**
-         * Initializes [TestFormula] instance with [initialOutput].
-         */
-        operator fun <Input, Output> invoke(
-            initialOutput: Output,
-            key: (Input) -> Any? = { null }
-        ): TestFormula<Input, Output> {
-            return object : TestFormula<Input, Output>() {
-                override fun initialOutput(): Output = initialOutput
-
-                override fun key(input: Input): Any? = key(input)
-            }
-        }
-    }
+class TestFormula<Input, Output> internal constructor(
+    private val parentFormula: IFormula<Input, Output>,
+    private val initialOutput: Output,
+) : Formula<Input, TestFormula.State<Output>, Output>() {
 
     data class State<Output>(
         val uniqueIdentifier: Long,
-        val key: Any?,
         val output: Output
     )
 
-    data class Value<Input, Output>(
-        val key: Any?,
-        val input: Input,
-        val output: Output,
-        val onNewOutput: (Output) -> Unit
-    )
-
     private val identifierGenerator = AtomicLong(0)
-
-    /**
-     * Uses initial input as key (to be decided if its robust enough)
-     */
-    private val stateMap = mutableMapOf<Any?, Value<Input, Output>>()
-
-    abstract fun initialOutput(): Output
+    private val runningInstanceManager = RunningInstanceManager<Input, Output>(
+        formulaKeyFactory = { parentFormula.key(it) },
+    )
 
     override fun initialState(input: Input): State<Output> {
         return State(
             uniqueIdentifier = identifierGenerator.getAndIncrement(),
-            key = key(input),
-            output = initialOutput(),
+            output = initialOutput,
         )
     }
 
     override fun Snapshot<Input, State<Output>>.evaluate(): Evaluation<Output> {
-        stateMap[state.uniqueIdentifier] = Value(
-            key = state.key,
+        runningInstanceManager.onEvaluate(
+            uniqueIdentifier = state.uniqueIdentifier,
             input = input,
             output = state.output,
             onNewOutput = context.onEvent {
-                transition(state.copy(output = it))
+                val newState = state.copy(output = it)
+                transition(newState)
             }
         )
 
@@ -73,7 +77,7 @@ abstract class TestFormula<Input, Output> :
             output = state.output,
             actions = context.actions {
                 Action.onTerminate().onEvent {
-                    stateMap.remove(state.uniqueIdentifier)
+                    runningInstanceManager.onTerminate(state.uniqueIdentifier)
                     none()
                 }
             }
@@ -84,23 +88,23 @@ abstract class TestFormula<Input, Output> :
      * Emits a new [Output].
      */
     fun output(output: Output) {
-        val update = getMostRecentRunningFormula().onNewOutput
+        val update = runningInstanceManager.mostRecentInstance().onNewOutput
         update(output)
     }
 
     fun output(key: Any?, output: Output) {
-        val instance = getRunningFormulaByKey(key)
+        val instance = runningInstanceManager.instanceByKey(key)
         instance.onNewOutput(output)
     }
 
     fun updateOutput(modify: Output.() -> Output) {
-        val formulaValue = getMostRecentRunningFormula()
+        val formulaValue = runningInstanceManager.mostRecentInstance()
         val newOutput = formulaValue.output.modify()
         formulaValue.onNewOutput(newOutput)
     }
 
     fun updateOutput(key: Any?, modify: Output.() -> Output) {
-        val formulaValue = getRunningFormulaByKey(key)
+        val formulaValue = runningInstanceManager.instanceByKey(key)
         val newOutput = formulaValue.output.modify()
         formulaValue.onNewOutput(newOutput)
     }
@@ -109,35 +113,25 @@ abstract class TestFormula<Input, Output> :
      * Performs an interaction on the current [Input] passed by the parent.
      */
     fun input(interact: Input.() -> Unit) {
-        val input = getMostRecentRunningFormula().input
-        interact(input)
+        lastInput().interact()
     }
 
     /**
      * Performs an interaction on the current [Input] passed by the parent.
      */
     fun input(key: Any?, interact: Input.() -> Unit) {
-        val instance = getRunningFormulaByKey(key)
-        instance.input.interact()
+        lastInputByKey(key).interact()
     }
 
     fun assertRunningCount(expected: Int) {
-        val count = stateMap.size
-        if (count != expected) {
-            throw AssertionError("Expected $expected running formulas, but there were $count instead")
-        }
+        runningInstanceManager.assertRunningCount(expected)
     }
 
-    private fun getMostRecentRunningFormula(): Value<Input, Output> {
-        return requireNotNull(stateMap.values.lastOrNull()) {
-            "Formula is not running"
-        }
+    fun lastInput(): Input {
+        return runningInstanceManager.mostRecentInstance().inputs.last()
     }
 
-    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"
-        }
+    fun lastInputByKey(key: Any?): Input {
+        return runningInstanceManager.instanceByKey(key).inputs.last()
     }
 }
\ No newline at end of file
diff --git a/formula-test/src/main/java/com/instacart/formula/test/utils/RunningInstanceManager.kt b/formula-test/src/main/java/com/instacart/formula/test/utils/RunningInstanceManager.kt
new file mode 100644
index 00000000..7363b23b
--- /dev/null
+++ b/formula-test/src/main/java/com/instacart/formula/test/utils/RunningInstanceManager.kt
@@ -0,0 +1,62 @@
+package com.instacart.formula.test.utils
+
+internal class RunningInstanceManager<Input, Output>(
+    private val formulaKeyFactory: (Input) -> Any?,
+) {
+    data class State<Input, Output>(
+        val key: Any?,
+        val inputs: List<Input>,
+        val output: Output,
+        val onNewOutput: (Output) -> Unit,
+    )
+
+    // Uses identifier as key that is generated within [initialState]
+    private val runningInstanceStates = mutableMapOf<Long, State<Input, Output>>()
+
+    fun onEvaluate(
+        uniqueIdentifier: Long,
+        input: Input,
+        output: Output,
+        onNewOutput: (Output) -> Unit,
+    ) {
+        val currentValue = runningInstanceStates[uniqueIdentifier]
+        val newInputList = if (currentValue == null) {
+            listOf(input)
+        } else if (currentValue.inputs.last() == input) {
+            currentValue.inputs
+        } else {
+            currentValue.inputs + input
+        }
+
+        runningInstanceStates[uniqueIdentifier] = State(
+            key = formulaKeyFactory(input),
+            inputs = newInputList,
+            output = output,
+            onNewOutput = onNewOutput,
+        )
+    }
+
+    fun onTerminate(uniqueIdentifier: Long) {
+        runningInstanceStates.remove(uniqueIdentifier)
+    }
+
+    fun mostRecentInstance(): State<Input, Output> {
+        return requireNotNull(runningInstanceStates.values.lastOrNull()) {
+            "Formula is not running"
+        }
+    }
+
+    fun instanceByKey(key: Any?): State<Input, Output> {
+        return requireNotNull(runningInstanceStates.entries.firstOrNull { it.value.key == key }?.value) {
+            val existingKeys = runningInstanceStates.entries.map { it.value.key }
+            "Formula for $key is not running, there are $existingKeys running"
+        }
+    }
+
+    fun assertRunningCount(expected: Int) {
+        val count = runningInstanceStates.size
+        if (count != expected) {
+            throw AssertionError("Expected $expected running formulas, but there were $count instead")
+        }
+    }
+}
\ No newline at end of file
diff --git a/formula-test/src/test/java/com/instacart/formula/test/SimpleFormula.kt b/formula-test/src/test/java/com/instacart/formula/test/SimpleFormula.kt
index da24fde2..bf61d988 100644
--- a/formula-test/src/test/java/com/instacart/formula/test/SimpleFormula.kt
+++ b/formula-test/src/test/java/com/instacart/formula/test/SimpleFormula.kt
@@ -8,14 +8,27 @@ 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"
+    override fun key(input: Input): Any? {
+        return if (useCustomKey) {
+            CUSTOM_KEY
+        } else {
+            super.key(input)
+        }
+    }
+
+    abstract val useCustomKey: Boolean
+
+    companion object {
+        const val CUSTOM_KEY = "simple-formula-key"
+    }
 }
 
 class TestSimpleFormula(
     private val initialOutput: Output = Output(
         outputId = 0,
         text = "",
-    )
+    ),
+    override val useCustomKey: Boolean = true,
 ) : SimpleFormula {
     override val implementation = testFormula(
         initialOutput = initialOutput,
diff --git a/formula-test/src/test/java/com/instacart/formula/test/TestFormulaTest.kt b/formula-test/src/test/java/com/instacart/formula/test/TestFormulaTest.kt
index 7cd92f0d..7bcfa142 100644
--- a/formula-test/src/test/java/com/instacart/formula/test/TestFormulaTest.kt
+++ b/formula-test/src/test/java/com/instacart/formula/test/TestFormulaTest.kt
@@ -13,9 +13,10 @@ class TestFormulaTest {
     }
 
     @Test fun `emits initial output when subscribed`() {
-        val formula = TestSimpleFormula()
+        val initialOutput = SimpleFormula.Output(100, "random")
+        val formula = TestSimpleFormula(initialOutput)
         formula.test().input(SimpleFormula.Input()).output {
-            assertThat(this).isEqualTo(formula.implementation.initialOutput())
+            assertThat(this).isEqualTo(initialOutput)
         }
     }
 
@@ -135,4 +136,16 @@ class TestFormulaTest {
             assertThat(this).isEqualTo(myInput)
         }
     }
+
+    @Test fun `default key is null`() {
+        val formula = TestSimpleFormula(useCustomKey = false)
+        val key = formula.key(SimpleFormula.Input())
+        assertThat(key).isNull()
+    }
+
+    @Test fun `custom key is applied correctly`() {
+        val formula = TestSimpleFormula(useCustomKey = true)
+        val key = formula.key(SimpleFormula.Input())
+        assertThat(key).isEqualTo(SimpleFormula.CUSTOM_KEY)
+    }
 }
diff --git a/formula/src/main/java/com/instacart/formula/StatelessFormula.kt b/formula/src/main/java/com/instacart/formula/StatelessFormula.kt
index 0817e985..660feb90 100644
--- a/formula/src/main/java/com/instacart/formula/StatelessFormula.kt
+++ b/formula/src/main/java/com/instacart/formula/StatelessFormula.kt
@@ -37,13 +37,4 @@ abstract class StatelessFormula<Input, Output> : IFormula<Input, Output> {
      * All side-effects should happen as part of event listeners or [actions][Evaluation.actions].
      */
      abstract fun Snapshot<Input, Unit>.evaluate(): Evaluation<Output>
-
-    /**
-     * A unique identifier used to distinguish formulas of the same type.
-     *
-     * ```
-     * override fun key(input: ItemInput) = input.itemId
-     * ```
-     */
-    override fun key(input: Input): Any? = null
 }