diff --git a/goth/runner/cli/yagna_payment_cmd.py b/goth/runner/cli/yagna_payment_cmd.py index a78cc5de..85c384c7 100644 --- a/goth/runner/cli/yagna_payment_cmd.py +++ b/goth/runner/cli/yagna_payment_cmd.py @@ -127,3 +127,9 @@ def payment_drivers( args = make_args("payment", "drivers") output = self.run_json_command(Dict, *args) return {key: Driver.from_dict(val) for key, val in output.items()} + + def payment_release_allocations(self: CommandRunner) -> None: + """Run ` payment release-allocations` without any extra args.""" + + args = make_args("payment", "release-allocations") + self.run_command(*args) diff --git a/goth/runner/probe/mixin.py b/goth/runner/probe/mixin.py index e3d312e9..7dcf658f 100644 --- a/goth/runner/probe/mixin.py +++ b/goth/runner/probe/mixin.py @@ -259,8 +259,8 @@ async def pay_invoices( timestamp=datetime.now(timezone.utc), payment_platform=self.payment_config.platform_string, ) - allocation_result = await self.api.payment.create_allocation(allocation) - logger.debug("Created allocation. id=%s", allocation_result) + + allocation_result = await self._create_allocation(allocation) acceptance = Acceptance( total_amount_accepted=invoice_event.amount, @@ -268,3 +268,41 @@ async def pay_invoices( ) await self.api.payment.accept_invoice(invoice_event.invoice_id, acceptance) logger.debug("Accepted invoice. id=%s", invoice_event.invoice_id) + + @step() + async def create_allocation( + self: ProbeProtocol, timeout: Optional[datetime] = None + ) -> Allocation: + """Call create_allocation on the market api.""" + + allocation = Allocation( + allocation_id="", + total_amount=0, + spent_amount=0, + remaining_amount=0, + make_deposit=True, + timestamp=datetime.now(timezone.utc), + timeout=timeout, + payment_platform=self.payment_config.platform_string, + ) + + allocation_result = await self._create_allocation(allocation) + + return allocation_result + + @step() + async def get_allocation(self: ProbeProtocol, allocation_id: str) -> Allocation: + """Call get_allocation on the market api.""" + + allocation_result = await self.api.payment.get_allocation(allocation_id) + + return allocation_result + + async def _create_allocation( + self: ProbeProtocol, allocation: Allocation + ) -> Allocation: + + allocation_result = await self.api.payment.create_allocation(allocation) + logger.debug("Created allocation. id=%s", allocation_result.allocation_id) + + return allocation_result