Increase code coverage (pt3).

formula/src/main/java/com/instacart/formula/internal/ListenerImpl.kt

@@ -8,14 +8,15 @@ import com.instacart.formula.plugin.Dispatcher
  * Note: this class is not a data class because equality is based on instance and not [key].
  */
 @PublishedApi
-internal class ListenerImpl<Input, State, EventT>(private val key: Any) : Listener<EventT> {
+internal class ListenerImpl<Input, State, EventT>(
+    private val key: Any,
+    private var transition: Transition<Input, State, EventT>
+) : Listener<EventT> {
 
     @Volatile private var manager: FormulaManagerImpl<Input, State, *>? = null
     @Volatile private var snapshotImpl: SnapshotImpl<Input, State>? = null
     @Volatile private var executionType: Transition.ExecutionType? = null
 
-    private lateinit var transition: Transition<Input, State, EventT>
-
     override fun invoke(event: EventT) {
         // TODO: log if null listener (it might be due to formula removal or due to callback removal)
         val manager = manager ?: return

formula/src/main/java/com/instacart/formula/internal/Listeners.kt

@@ -1,18 +1,24 @@
 package com.instacart.formula.internal
 
+import com.instacart.formula.Transition
+
 internal class Listeners {
     private var listeners: SingleRequestMap<Any, ListenerImpl<*, *, *>>? = null
     private var indexes: MutableMap<Any, Int>? = null
 
-    fun <Input, State, Event> initOrFindListener(key: Any, useIndex: Boolean): ListenerImpl<Input, State, Event> {
-        val currentHolder = listenerHolder<Input, State, Event>(key)
+    fun <Input, State, Event> initOrFindListener(
+        key: Any,
+        useIndex: Boolean,
+        transition: Transition<Input, State, Event>
+    ): ListenerImpl<Input, State, Event> {
+        val currentHolder = listenerHolder(key, transition)
         return if (!currentHolder.requested) {
             currentHolder.requested = true
             currentHolder.value as ListenerImpl<Input, State, Event>
         } else if (useIndex) {
             val index = nextIndex(key)
             val indexedKey = IndexedKey(key, index)
-            initOrFindListener(indexedKey, useIndex)
+            initOrFindListener(indexedKey, useIndex, transition)
         } else {
             throw IllegalStateException("Listener $key is already defined. Unexpected issue.")
         }
@@ -60,13 +66,16 @@ internal class Listeners {
         return index
     }
 
-    private fun <Input, State, Event> listenerHolder(key: Any): SingleRequestHolder<ListenerImpl<*, *, *>> {
+    private fun <Input, State, Event> listenerHolder(
+        key: Any,
+        transition: Transition<Input, State, Event>
+    ): SingleRequestHolder<ListenerImpl<*, *, *>> {
         val listeners = listeners ?: run {
             val initialized: SingleRequestMap<Any, ListenerImpl<*, *, *>> = mutableMapOf()
             this.listeners = initialized
             initialized
         }
 
-        return listeners.findOrInit(key) { ListenerImpl<Input, State, Event>(key) }
+        return listeners.findOrInit(key) { ListenerImpl(key, transition) }
     }
 }

formula/src/main/java/com/instacart/formula/internal/SnapshotImpl.kt

@@ -64,7 +64,7 @@ internal class SnapshotImpl<out Input, State>(
         transition: Transition<Input, State, Event>
     ): Listener<Event> {
         ensureNotRunning()
-        val listener = listeners.initOrFindListener<Input, State, Event>(key, useIndex)
+        val listener = listeners.initOrFindListener(key, useIndex, transition)
         listener.setDependencies(delegate, this, executionType, transition)
         return listener
     }

formula/src/test/java/com/instacart/formula/FormulaRuntimeTest.kt

@@ -663,6 +663,79 @@ class FormulaRuntimeTest(val runtime: TestableRuntime, val name: String) {
             }
     }
 
+    @Test
+    fun `dispatched event is ignored if listener was disabled before event is processed`() {
+        data class State(
+            val listenerEnabled: Boolean = true,
+            val value: Int = 0,
+        )
+
+        data class Output(
+            val value: Int,
+            val increment: () -> Unit,
+        )
+
+        val disableListenerRelay = runtime.newRelay()
+        val formula = object : Formula<Unit, State, Output>() {
+            override fun initialState(input: Unit): State = State()
+
+            override fun Snapshot<Unit, State>.evaluate(): Evaluation<Output> {
+                val increment = if (state.listenerEnabled) {
+                    context.callback {
+                        transition(state.copy(value = state.value.inc()))
+                    }
+                } else {
+                    context.callback { none() }
+                }
+                return Evaluation(
+                    output = Output(
+                        value = state.value,
+                        increment = increment
+                    ),
+                    actions = context.actions {
+                        disableListenerRelay.action().onEvent {
+                            val listener = state.copy(listenerEnabled = false)
+                            transition(listener)
+                        }
+                    }
+                )
+            }
+        }
+
+        val dispatcher = object : Dispatcher {
+            var dispatches = mutableListOf<() -> Unit>()
+
+            override fun isDispatchNeeded(): Boolean {
+                return true
+            }
+
+            override fun dispatch(executable: () -> Unit) {
+                dispatches.add(executable)
+            }
+
+            fun executeAndClear() {
+                val local = dispatches
+                dispatches = mutableListOf()
+                local.forEach { it.invoke() }
+            }
+        }
+        val observer = runtime.test(formula, defaultDispatcher = dispatcher)
+
+        // Initialize formula
+        observer.input(Unit)
+        dispatcher.executeAndClear()
+
+        // First
+        val increment = observer.values().last().increment
+        disableListenerRelay.triggerEvent()
+        increment()
+        increment()
+        increment()
+
+        dispatcher.executeAndClear()
+        observer.output { assertThat(value).isEqualTo(0) }
+    }
+
     @Test fun `dispatching does not affect event order`() {
         var observer: TestFormulaObserver<Unit, OptionalCallbackFormula.Output, OptionalCallbackFormula>? = null
         FormulaPlugins.setPlugin(object : Plugin {