import Foundation

class MockFunc<Input, Output> {
    typealias MockedFunc = (Input) -> (Output)
    typealias Callback = (Input) -> (Output)

    var calls: [(Input)] = []
    var callback: Callback!

    func call(input: Input) -> (Output) {
        calls.append(input)
        return callback(input)
    }

    func when(callback: Callback) {
        self.callback = callback
    }
}
import Foundation

class MockFunc<Input, Output> {
    typealias MockedFunc = (Input) -> (Output)
    var calls: [(Input)] = []
    var result: (Output)!

    func call(input: Input) -> (Output) {
        calls.append(input)
        return result
    }
}
--- 
+++ 
@@ -2,11 +2,17 @@
 
 class MockFunc<Input, Output> {
     typealias MockedFunc = (Input) -> (Output)
+    typealias Callback = (Input) -> (Output)
+
     var calls: [(Input)] = []
-    var result: (Output)!
+    var callback: Callback!
 
     func call(input: Input) -> (Output) {
         calls.append(input)
-        return result
+        return callback(input)
+    }
+
+    func when(callback: Callback) {
+        self.callback = callback
     }
 }
Add when() method to replace result field
