diff --git a/goth/runner/probe/mixin.py b/goth/runner/probe/mixin.py
index 7dcf658f..8448ad10 100644
--- a/goth/runner/probe/mixin.py
+++ b/goth/runner/probe/mixin.py
@@ -271,13 +271,13 @@ async def pay_invoices(
 
     @step()
     async def create_allocation(
-        self: ProbeProtocol, timeout: Optional[datetime] = None
+        self: ProbeProtocol, timeout: Optional[datetime] = None, total_amount=0
     ) -> Allocation:
         """Call create_allocation on the market api."""
 
         allocation = Allocation(
             allocation_id="",
-            total_amount=0,
+            total_amount=total_amount,
             spent_amount=0,
             remaining_amount=0,
             make_deposit=True,