From 8098f5220bda7dc53e13b5417a8e9cf74059c3c8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafa=C5=82=20Borysionek?=
 <22373844+MrDarthShoe@users.noreply.github.com>
Date: Fri, 5 Aug 2022 10:00:21 +0200
Subject: [PATCH 1/2] Allow total amount allocation

---
 goth/runner/probe/mixin.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/goth/runner/probe/mixin.py b/goth/runner/probe/mixin.py
index 7dcf658f..f5be6799 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,

From 2c38b939344c9fd649e37a75b94478f4c1df8ad6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafa=C5=82=20Borysionek?=
 <22373844+MrDarthShoe@users.noreply.github.com>
Date: Fri, 5 Aug 2022 10:43:30 +0200
Subject: [PATCH 2/2] black format

---
 goth/runner/probe/mixin.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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