added pytest timeout

agentcoinorg · Jun 28, 2024 · 0fa32df · 0fa32df
1 parent e95d49d
commit 0fa32df
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 1 deletion.
diff --git a/autotx/tests/agents/token/send/test_send.py b/autotx/tests/agents/token/send/test_send.py
@@ -1,6 +1,9 @@
+import pytest
+
 from autotx.utils.ethereum import get_erc20_balance
 from autotx.utils.ethereum.get_native_balance import get_native_balance
 
+@pytest.mark.timeout(60)
 def test_send_native(smart_account, auto_tx, test_accounts):
     receiver = test_accounts[0]
     balance = get_native_balance(smart_account.web3, receiver)
@@ -11,6 +14,7 @@ def test_send_native(smart_account, auto_tx, test_accounts):
     balance = get_native_balance(smart_account.web3, receiver)
     assert balance == 1
 
+@pytest.mark.timeout(60)
 def test_send_erc20(smart_account, auto_tx, usdc, test_accounts):
 
     receiver = test_accounts[0]
@@ -25,6 +29,7 @@ def test_send_erc20(smart_account, auto_tx, usdc, test_accounts):
 
     assert balance + 10 == new_balance
 
+@pytest.mark.timeout(60)
 def test_send_native_sequential(smart_account, auto_tx, test_accounts):
     receiver = test_accounts[0]
 
@@ -38,6 +43,7 @@ def test_send_native_sequential(smart_account, auto_tx, test_accounts):
     balance = get_native_balance(smart_account.web3, receiver)
     assert balance == 1.5
 
+@pytest.mark.timeout(60)
 def test_send_erc20_parallel(smart_account, auto_tx, usdc, test_accounts):
 
     receiver_one = test_accounts[0]
@@ -56,6 +62,7 @@ def test_send_erc20_parallel(smart_account, auto_tx, usdc, test_accounts):
     assert balance_one + 2 == new_balance_one
     assert balance_two + 3 == new_balance_two
 
+@pytest.mark.timeout(60)
 def test_send_eth_multiple(smart_account, auto_tx, usdc, test_accounts):
 
     receiver_1 = test_accounts[0]

diff --git a/autotx/tests/agents/token/test_swap.py b/autotx/tests/agents/token/test_swap.py
@@ -1,9 +1,12 @@
+import pytest
+
 from autotx.utils.ethereum.get_erc20_balance import get_erc20_balance
 from autotx.utils.ethereum.networks import NetworkInfo
 from autotx.eth_address import ETHAddress
 
 DIFFERENCE_PERCENTAGE = 1.01
 
+@pytest.mark.timeout(60)
 def test_swap_with_non_default_token(smart_account, auto_tx):
     network_info = NetworkInfo(smart_account.web3.eth.chain_id)
     shib_address = ETHAddress(network_info.tokens["shib"])
@@ -18,6 +21,7 @@ def test_swap_with_non_default_token(smart_account, auto_tx):
     expected_shib_amount = 100000
     assert expected_shib_amount <= new_balance <= expected_shib_amount * DIFFERENCE_PERCENTAGE
 
+@pytest.mark.timeout(60)
 def test_swap_native(smart_account, auto_tx):
     network_info = NetworkInfo(smart_account.web3.eth.chain_id)
     usdc_address = ETHAddress(network_info.tokens["usdc"])
@@ -29,6 +33,7 @@ def test_swap_native(smart_account, auto_tx):
     expected_usdc_amount = 100
     assert expected_usdc_amount <= new_balance <= expected_usdc_amount * DIFFERENCE_PERCENTAGE
 
+@pytest.mark.timeout(60)
 def test_swap_multiple_1(smart_account, auto_tx):
     network_info = NetworkInfo(smart_account.web3.eth.chain_id)
     usdc_address = ETHAddress(network_info.tokens["usdc"])
@@ -46,6 +51,7 @@ def test_swap_multiple_1(smart_account, auto_tx):
     assert expected_usdc_amount <= usdc_balance <= expected_usdc_amount_plus_slippage - expected_usdc_amount
     assert wbtc_balance < get_erc20_balance(smart_account.web3, wbtc_address, smart_account.address)
 
+@pytest.mark.timeout(60)
 def test_swap_multiple_2(smart_account, auto_tx):
     network_info = NetworkInfo(smart_account.web3.eth.chain_id)
     usdc_address = ETHAddress(network_info.tokens["usdc"])
@@ -61,6 +67,7 @@ def test_swap_multiple_2(smart_account, auto_tx):
     assert expected_amount <= usdc_balance
     assert wbtc_balance < get_erc20_balance(smart_account.web3, wbtc_address, smart_account.address)
 
+@pytest.mark.timeout(60)
 def test_swap_triple(smart_account, auto_tx): 
     network_info = NetworkInfo(smart_account.web3.eth.chain_id)
     usdc_address = ETHAddress(network_info.tokens["usdc"])
@@ -81,6 +88,7 @@ def test_swap_triple(smart_account, auto_tx):
     assert expected_uni_amount <= uni_balance <= expected_uni_amount * DIFFERENCE_PERCENTAGE
     assert expected_wbtc_amount <= wbtc_balance <= expected_wbtc_amount * DIFFERENCE_PERCENTAGE
 
+@pytest.mark.timeout(60)
 def test_swap_complex_1(smart_account, auto_tx): # This one is complex because it confuses the LLM with WBTC amount
     network_info = NetworkInfo(smart_account.web3.eth.chain_id)
     usdc_address = ETHAddress(network_info.tokens["usdc"])
@@ -95,6 +103,7 @@ def test_swap_complex_1(smart_account, auto_tx): # This one is complex because i
     assert expected_usdc_amount <= usdc_balance <= expected_usdc_amount * DIFFERENCE_PERCENTAGE
     assert wbtc_balance < get_erc20_balance(smart_account.web3, wbtc_address, smart_account.address)
 
+@pytest.mark.timeout(60)
 def test_swap_complex_2(smart_account, auto_tx): # This one is complex because it confuses the LLM with WBTC amount
     network_info = NetworkInfo(smart_account.web3.eth.chain_id)
     usdc_address = ETHAddress(network_info.tokens["usdc"])

diff --git a/autotx/tests/agents/token/test_swap_and_send.py b/autotx/tests/agents/token/test_swap_and_send.py
@@ -1,9 +1,11 @@
+import pytest
 from autotx.utils.ethereum import get_erc20_balance, get_native_balance
 from autotx.utils.ethereum.networks import NetworkInfo
 from autotx.eth_address import ETHAddress
 
 DIFFERENCE_PERCENTAGE = 1.01
 
+@pytest.mark.timeout(60)
 def test_swap_and_send_simple(smart_account, auto_tx, test_accounts):
     network_info = NetworkInfo(smart_account.web3.eth.chain_id)
     wbtc_address = ETHAddress(network_info.tokens["wbtc"])
@@ -20,6 +22,7 @@ def test_swap_and_send_simple(smart_account, auto_tx, test_accounts):
     assert excepted_safe_wbtc_balance <= new_wbtc_safe_address <= new_wbtc_safe_address * DIFFERENCE_PERCENTAGE
     assert new_receiver_wbtc_balance == 0.01
 
+@pytest.mark.timeout(60)
 def test_swap_and_send_complex(smart_account, auto_tx, test_accounts):
     network_info = NetworkInfo(smart_account.web3.eth.chain_id)
     usdc_address = ETHAddress(network_info.tokens["usdc"])
@@ -41,6 +44,7 @@ def test_swap_and_send_complex(smart_account, auto_tx, test_accounts):
     assert expected_usdc_safe_balance <= new_usdc_safe_address <= expected_usdc_safe_balance * DIFFERENCE_PERCENTAGE
     assert new_receiver_usdc_balance == 50
 
+@pytest.mark.timeout(60)
 def test_send_and_swap_simple(smart_account, auto_tx, test_accounts):
     network_info = NetworkInfo(smart_account.web3.eth.chain_id)
     wbtc_address = ETHAddress(network_info.tokens["wbtc"])
@@ -64,6 +68,7 @@ def test_send_and_swap_simple(smart_account, auto_tx, test_accounts):
     assert new_receiver_wbtc_balance == receiver_wbtc_balance
     assert new_receiver_native_balance == receiver_native_balance + 0.1
 
+@pytest.mark.timeout(60)
 def test_send_and_swap_complex(smart_account, auto_tx, test_accounts):
     network_info = NetworkInfo(smart_account.web3.eth.chain_id)
     usdc_address = ETHAddress(network_info.tokens["usdc"])

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -20,6 +20,7 @@ uvicorn = "^0.29.0"
 supabase = "^2.5.0"
 llama-cpp-python = "^0.2.78"
 aiohttp = "^3.9.5"
+pytest-timeout = "^2.3.1"
 
 [tool.poetry.group.dev.dependencies]
 mypy = "^1.8.0"