diff --git a/Source/TestWait/TestMock/WaitMockFunction.swift b/Source/TestWait/TestMock/WaitMockFunction.swift
index ad03a6d..9bbc5d8 100644
--- a/Source/TestWait/TestMock/WaitMockFunction.swift
+++ b/Source/TestWait/TestMock/WaitMockFunction.swift
@@ -29,9 +29,10 @@ public struct WaitMockFunction {
             if until(argument, result) {
                 callArgument = argument
                 callResult = result
-                callback(argument, result)
                 semaphore.signal()
+                mock.callback = callback
             }
+            callback(argument, result)
         }
         try block()