Add manage_flow to superfluid.py

Author: philogicae

src/aleph/sdk/chains/ethereum.py

@@ -22,12 +22,13 @@
     BALANCEOF_ABI,
     MIN_ETH_BALANCE,
     MIN_ETH_BALANCE_WEI,
+    FlowUpdate,
+    from_wei_token,
     get_chain_id,
     get_chains_with_super_token,
     get_rpc,
     get_super_token_address,
     get_token_address,
-    to_human_readable_token,
 )
 from ..exceptions import BadSignatureError
 from ..utils import bytes_from_hex
@@ -107,7 +108,7 @@ def can_transact(self, block=True) -> bool:
         if not valid and block:
             raise InsufficientFundsError(
                 required_funds=MIN_ETH_BALANCE,
-                available_funds=to_human_readable_token(balance),
+                available_funds=float(from_wei_token(balance)),
             )
         return valid
 
@@ -162,32 +163,55 @@ def get_super_token_balance(self) -> Decimal:
                 return Decimal(contract.functions.balanceOf(self.get_address()).call())
         return Decimal(0)
 
+    @property
+    def has_superfluid_connector(self) -> bool:
+        return self.superfluid_connector is not None
+
+    def can_start_flow(self, flow: Decimal) -> Awaitable[bool]:
+        """Check if the account has enough funds to start a Superfluid flow of the given size."""
+        if not self.has_superfluid_connector:
+            raise ValueError("Superfluid connector is required to check a flow")
+        return self.superfluid_connector.can_start_flow(flow)
+
     def create_flow(self, receiver: str, flow: Decimal) -> Awaitable[str]:
         """Creat a Superfluid flow between this account and the receiver address."""
-        if not self.superfluid_connector:
+        if not self.has_superfluid_connector:
             raise ValueError("Superfluid connector is required to create a flow")
         return self.superfluid_connector.create_flow(receiver=receiver, flow=flow)
 
     def get_flow(self, receiver: str) -> Awaitable[Web3FlowInfo]:
         """Get the Superfluid flow between this account and the receiver address."""
-        if not self.superfluid_connector:
+        if not self.has_superfluid_connector:
             raise ValueError("Superfluid connector is required to get a flow")
         return self.superfluid_connector.get_flow(
             sender=self.get_address(), receiver=receiver
         )
 
     def update_flow(self, receiver: str, flow: Decimal) -> Awaitable[str]:
         """Update the Superfluid flow between this account and the receiver address."""
-        if not self.superfluid_connector:
+        if not self.has_superfluid_connector:
             raise ValueError("Superfluid connector is required to update a flow")
         return self.superfluid_connector.update_flow(receiver=receiver, flow=flow)
 
     def delete_flow(self, receiver: str) -> Awaitable[str]:
         """Delete the Superfluid flow between this account and the receiver address."""
-        if not self.superfluid_connector:
+        if not self.has_superfluid_connector:
             raise ValueError("Superfluid connector is required to delete a flow")
         return self.superfluid_connector.delete_flow(receiver=receiver)
 
+    def manage_flow(
+        self,
+        receiver: str,
+        flow: Decimal,
+        update_type: FlowUpdate,
+    ) -> Awaitable[Optional[str]]:
+        """Manage the Superfluid flow between this account and the receiver address."""
+        if not self.has_superfluid_connector:
+            raise ValueError("Superfluid connector is required to manage a flow")
+        return self.superfluid_connector.manage_flow(
+            receiver=receiver, flow=flow, update_type=update_type
+        )
+
 
 def get_fallback_account(
     path: Optional[Path] = None, chain: Optional[Chain] = None

src/aleph/sdk/chains/evm.py

@@ -5,6 +5,7 @@
 from aleph_message.models import Chain
 from eth_account import Account  # type: ignore
 
+from ..evm_utils import FlowUpdate
 from .common import get_fallback_private_key
 from .ethereum import ETHAccount
 
@@ -29,6 +30,9 @@ def get_token_balance(self) -> Decimal:
     def get_super_token_balance(self) -> Decimal:
         raise ValueError(f"Super token not implemented for this chain {self.CHAIN}")
 
+    def can_start_flow(self, flow: Decimal) -> Awaitable[bool]:
+        raise ValueError(f"Flow checking not implemented for this chain {self.CHAIN}")
+
     def create_flow(self, receiver: str, flow: Decimal) -> Awaitable[str]:
         raise ValueError(f"Flow creation not implemented for this chain {self.CHAIN}")
 
@@ -41,6 +45,11 @@ def update_flow(self, receiver: str, flow: Decimal) -> Awaitable[str]:
     def delete_flow(self, receiver: str) -> Awaitable[str]:
         raise ValueError(f"Flow deletion not implemented for this chain {self.CHAIN}")
 
+    def manage_flow(
+        self, receiver: str, flow: Decimal, update_type: FlowUpdate
+    ) -> Awaitable[Optional[str]]:
+        raise ValueError(f"Flow management not implemented for this chain {self.CHAIN}")
+
 
 def get_fallback_account(
     path: Optional[Path] = None, chain: Optional[Chain] = None

src/aleph/sdk/connectors/superfluid.py

@@ -1,14 +1,19 @@
 from __future__ import annotations
 
 from decimal import Decimal
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 from eth_utils import to_normalized_address
 from superfluid import CFA_V1, Operation, Web3FlowInfo
 
 from aleph.sdk.exceptions import InsufficientFundsError
 
-from ..evm_utils import get_super_token_address, to_human_readable_token, to_wei_token
+from ..evm_utils import (
+    FlowUpdate,
+    from_wei_token,
+    get_super_token_address,
+    to_wei_token,
+)
 
 if TYPE_CHECKING:
     from aleph.sdk.chains.ethereum import ETHAccount
@@ -44,6 +49,7 @@ async def _execute_operation_with_account(self, operation: Operation) -> str:
         return await self.account._sign_and_send_transaction(populated_transaction)
 
     def can_start_flow(self, flow: Decimal, block=True) -> bool:
+        """Check if the account has enough funds to start a Superfluid flow of the given size."""
         valid = False
         if self.account.can_transact(block=block):
             balance = self.account.get_super_token_balance()
@@ -52,7 +58,7 @@ def can_start_flow(self, flow: Decimal, block=True) -> bool:
             if not valid and block:
                 raise InsufficientFundsError(
                     required_funds=float(MIN_FLOW_4H),
-                    available_funds=to_human_readable_token(balance),
+                    available_funds=float(from_wei_token(balance)),
                 )
         return valid
 
@@ -96,3 +102,51 @@ async def update_flow(self, receiver: str, flow: Decimal) -> str:
                 flow_rate=int(to_wei_token(flow)),
             ),
         )
+
+    async def manage_flow(
+        self,
+        receiver: str,
+        flow: Decimal,
+        update_type: FlowUpdate,
+    ) -> Optional[str]:
+        """
+        Update the flow of a Superfluid stream between a sender and receiver.
+        This function either increases or decreases the flow rate between the sender and receiver,
+        based on the update_type. If no flow exists and the update type is augmentation, it creates a new flow
+        with the specified rate. If the update type is reduction and the reduction amount brings the flow to zero
+        or below, the flow is deleted.
+
+        :param receiver: Address of the receiver in hexadecimal format.
+        :param flow: The flow rate to be added or removed (in ether).
+        :param update_type: The type of update to perform (augmentation or reduction).
+        :return: The transaction hash of the executed operation (create, update, or delete flow).
+        """
+
+        # Retrieve current flow info
+        flow_info: Web3FlowInfo = await self.account.get_flow(receiver)
+
+        current_flow_rate_wei: Decimal = Decimal(flow_info["flowRate"] or "0")
+        flow_rate_wei: int = int(to_wei_token(flow))
+
+        if update_type == FlowUpdate.INCREASE:
+            if current_flow_rate_wei > 0:
+                # Update existing flow by increasing the rate
+                new_flow_rate_wei = current_flow_rate_wei + flow_rate_wei
+                new_flow_rate_ether = from_wei_token(new_flow_rate_wei)
+                return await self.account.update_flow(receiver, new_flow_rate_ether)
+            else:
+                # Create a new flow if none exists
+                return await self.account.create_flow(receiver, flow)
+        else:
+            if current_flow_rate_wei > 0:
+                # Reduce the existing flow
+                new_flow_rate_wei = current_flow_rate_wei - flow_rate_wei
+                # Ensure to not leave infinitesimal flows
+                # Often, there were 1-10 wei remaining in the flow rate, which prevented the flow from being deleted
+                if new_flow_rate_wei > 99:
+                    new_flow_rate_ether = from_wei_token(new_flow_rate_wei)
+                    return await self.account.update_flow(receiver, new_flow_rate_ether)
+                else:
+                    # Delete the flow if the new flow rate is zero or negative
+                    return await self.account.delete_flow(receiver)
+        return None

src/aleph/sdk/evm_utils.py

@@ -1,4 +1,5 @@
 from decimal import Decimal
+from enum import Enum
 from typing import List, Optional, Union
 
 from aleph_message.models import Chain
@@ -21,11 +22,18 @@
 }]"""
 
 
-def to_human_readable_token(amount: Decimal) -> float:
-    return float(amount / (Decimal(10) ** Decimal(settings.TOKEN_DECIMALS)))
+class FlowUpdate(str, Enum):
+    REDUCE = "reduce"
+    INCREASE = "increase"
+
+
+def from_wei_token(amount: Decimal) -> Decimal:
+    """Converts the given wei value to ether."""
+    return amount / Decimal(10) ** Decimal(settings.TOKEN_DECIMALS)
 
 
 def to_wei_token(amount: Decimal) -> Decimal:
+    """Converts the given ether value to wei."""
     return amount * Decimal(10) ** Decimal(settings.TOKEN_DECIMALS)
 
 

tests/unit/test_superfluid.py

@@ -7,6 +7,7 @@
 from eth_utils import to_checksum_address
 
 from aleph.sdk.chains.ethereum import ETHAccount
+from aleph.sdk.evm_utils import FlowUpdate
 
 
 def generate_fake_eth_address():
@@ -24,6 +25,7 @@ def mock_superfluid():
         mock_superfluid.create_flow = AsyncMock(return_value="0xTransactionHash")
         mock_superfluid.delete_flow = AsyncMock(return_value="0xTransactionHash")
         mock_superfluid.update_flow = AsyncMock(return_value="0xTransactionHash")
+        mock_superfluid.manage_flow = AsyncMock(return_value="0xTransactionHash")
 
         # Mock get_flow to return a mock Web3FlowInfo
         mock_flow_info = {"timestamp": 0, "flowRate": 0, "deposit": 0, "owedDeposit": 0}
@@ -98,3 +100,14 @@ async def test_get_flow(eth_account, mock_superfluid):
     assert flow_info["flowRate"] == 0
     assert flow_info["deposit"] == 0
     assert flow_info["owedDeposit"] == 0
+
+
+@pytest.mark.asyncio
+async def test_manage_flow(eth_account, mock_superfluid):
+    receiver = generate_fake_eth_address()
+    flow = Decimal("0.005")
+
+    tx_hash = await eth_account.manage_flow(receiver, flow, FlowUpdate.INCREASE)
+
+    assert tx_hash == "0xTransactionHash"
+    mock_superfluid.manage_flow.assert_awaited_once()