expand produce customisation

Sources/StoryFlow/Interfaces/Helpers/Testing.swift

@@ -11,12 +11,32 @@ extension InputRequiring where Self: UIViewController {
 
 extension OutputProducing where Self: UIViewController {
 
-    public init(nibName: String? = nil, bundle: Bundle? = nil, produce: @escaping (OutputType) -> ()) {
+    /// - Parameters:
+    ///   - produceStub: Closure that optionally overrides default produce behavior.
+    ///   When nil is returned default produce behaviour will not be triggered.
+    public init(
+        nibName: String? = nil,
+        bundle: Bundle? = nil,
+        produceStub: @escaping (OutputType) -> OutputType?
+    ) {
         self.init(nibName: nibName, bundle: bundle)
-        produceStub = produce
+        self.produceStub = produceStub
+    }
+
+    /// - Parameters:
+    ///   - produceStub: Closure that overrides default produce behavior.
+    public init(
+        nibName: String? = nil,
+        bundle: Bundle? = nil,
+        produceStub: @escaping (OutputType) -> ()
+    ) {
+        self.init(nibName: nibName, bundle: bundle) {
+            produceStub($0)
+            return nil
+        }
     }
 
-    var produceStub: ((OutputType) -> ())? {
+    var produceStub: ((OutputType) -> OutputType?)? {
         get { return associated(with: &produceStubKey) }
         set { associate(newValue, with: &produceStubKey) }
     }
@@ -25,10 +45,33 @@ private var produceStubKey = 0
 
 extension InputRequiring where Self: UIViewController & OutputProducing {
 
-    public init(nibName: String? = nil, bundle: Bundle? = nil,
-                input: InputType, produce: @escaping (OutputType) -> ()) {
+    /// - Parameters:
+    ///   - produceStub: Closure that optionally overrides default produce behavior.
+    ///   When nil is returned default produce behaviour will not be triggered.
+    public init(
+        nibName: String? = nil,
+        bundle: Bundle? = nil,
+        input: InputType,
+        produceStub: @escaping (OutputType) -> OutputType?
+    ) {
         self.init(nibName: nibName, bundle: bundle)
         self.input = input
-        self.produceStub = produce
+        self.produceStub = produceStub
+    }
+
+    /// - Parameters:
+    ///   - produceStub: Closure that overrides default produce behavior.
+    public init(
+        nibName: String? = nil,
+        bundle: Bundle? = nil,
+        input: InputType,
+        produceStub: @escaping (OutputType) -> ()
+    ) {
+        self.init(nibName: nibName, bundle: bundle) {
+            produceStub($0)
+            return nil
+        }
+        self.input = input
     }
+
 }

Sources/StoryFlow/Interfaces/OuputProducing.swift

@@ -19,10 +19,17 @@ extension OutputProducing where Self: UIViewController {
      - Parameter output: The value being passed to next view controller.
      */
     public func produce(_ output: OutputType) {
-
-        if let produce = produceStub {
-            produce(output)
-        } else if let flow = flow {
+        if let produceStub {
+            if let customOutput = produceStub(output) {
+                _produce(customOutput)
+            }
+        } else {
+            _produce(output)
+        }
+    }
+
+    private func _produce(_ output: OutputType) {
+        if let flow = flow {
             flow.proceed(with: output, from: self)
         } else {
             implicitFlow.proceed(with: output, from: self)

StoryFlowTests/TestabilityTests.swift

@@ -25,7 +25,7 @@ class TestabilityTests: XCTestCase {
         class Vc: UIViewController, OutputProducing { typealias OutputType = T }
 
         var producedOutput: T!
-        let vc = Vc(produce: { producedOutput = $0 })
+        let vc = Vc(produceStub: { producedOutput = $0 })
 
         let output = T()
 
@@ -48,7 +48,7 @@ class TestabilityTests: XCTestCase {
         var producedOutput: T!
 
         // Act
-        let vc = Vc(input: value, produce: { producedOutput = $0 })
+        let vc = Vc(input: value, produceStub: { producedOutput = $0 })
         vc.produce(value)
 
         // Assert