Fix instance tests after changes

aleph-im · Jan 31, 2025 · 9223c80 · 9223c80
1 parent cae2281
commit 9223c80
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 10 deletions.
diff --git a/tests/unit/mocks.py b/tests/unit/mocks.py
@@ -3,7 +3,6 @@
 
 from aleph.sdk.chains.evm import EVMAccount
 from aleph.sdk.conf import settings
-from eth_utils.currency import to_wei
 from pydantic import BaseModel
 
 # Change to Aleph testnet
@@ -37,9 +36,6 @@ def create_mock_load_account():
     mock_loader = MagicMock(return_value=mock_account)
     mock_loader.return_value.get_super_token_balance = MagicMock(return_value=Decimal(10000 * (10**18)))
     mock_loader.return_value.can_transact = MagicMock(return_value=True)
-    mock_loader.return_value.superfluid_connector = MagicMock(can_start_flow=MagicMock(return_value=True))
-    mock_loader.return_value.get_flow = AsyncMock(return_value={"flowRate": to_wei(0.0001, unit="ether")})
-    mock_loader.return_value.create_flow = AsyncMock(return_value=FAKE_FLOW_HASH)
-    mock_loader.return_value.update_flow = AsyncMock(return_value=FAKE_FLOW_HASH)
-    mock_loader.return_value.delete_flow = AsyncMock(return_value=FAKE_FLOW_HASH)
+    mock_loader.return_value.can_start_flow = MagicMock(return_value=True)
+    mock_loader.return_value.manage_flow = AsyncMock(return_value=FAKE_FLOW_HASH)
     return mock_loader
diff --git a/tests/unit/test_instance.py b/tests/unit/test_instance.py
@@ -427,7 +427,6 @@ async def test_create_instance(args, expected):
     mock_crn_info = create_mock_crn_info()
     mock_validated_int_prompt = MagicMock(return_value=1)
     mock_wait_for_processed_instance = AsyncMock()
-    mock_update_flow = AsyncMock(return_value="fake_flow_hash")
     mock_wait_for_confirmed_flow = AsyncMock()
 
     @patch("aleph_client.commands.instance.validate_ssh_pubkey_file", mock_validate_ssh_pubkey_file)
@@ -437,7 +436,6 @@ async def test_create_instance(args, expected):
     @patch("aleph_client.commands.instance.CRNInfo", mock_crn_info)
     @patch("aleph_client.commands.instance.validated_int_prompt", mock_validated_int_prompt)
     @patch("aleph_client.commands.instance.wait_for_processed_instance", mock_wait_for_processed_instance)
-    @patch("aleph_client.commands.instance.update_flow", mock_update_flow)
     @patch("aleph_client.commands.instance.wait_for_confirmed_flow", mock_wait_for_confirmed_flow)
     @patch("aleph_client.commands.instance.VmClient", mock_vm_client_class)
     async def create_instance(instance_spec):
@@ -474,7 +472,7 @@ async def create_instance(instance_spec):
     mock_auth_client.create_instance.assert_called_once()
     if args["payment_type"] == "superfluid":
         mock_wait_for_processed_instance.assert_called_once()
-        mock_update_flow.assert_called_once()
+        mock_account.manage_flow.assert_called_once()
         mock_wait_for_confirmed_flow.assert_called_once()
         mock_vm_client.start_instance.assert_called_once()
     assert returned == expected
@@ -531,7 +529,7 @@ async def delete_instance():
         )
         mock_auth_client.get_message.assert_called_once()
         mock_vm_client.erase_instance.assert_called_once()
-        mock_account.delete_flow.assert_awaited_once()
+        mock_account.manage_flow.assert_awaited_once()
         mock_auth_client.forget.assert_called_once()
 
     await delete_instance()