diff --git a/ui_workflows/base.py b/ui_workflows/base.py index 19d73761..8642e56f 100644 --- a/ui_workflows/base.py +++ b/ui_workflows/base.py @@ -34,6 +34,15 @@ class Result: is_approval_tx: bool = False # NOTE: Field deprecated, use Multi-step workflow approach parsed_user_request: str = '' # NOTE: Field deprecated, use Multi-step workflow approach +@dataclass +class SingleStepResult: + status: Literal['success', 'error', 'terminated'] + workflow_type: str + description: str + user_action_type: Literal['tx', 'acknowledge'] + tx: any = None + error_msg: Optional[str] = None + @dataclass class MultiStepResult: status: Literal['success', 'error', 'terminated'] @@ -238,6 +247,60 @@ def _intercept_rpc_node_reqs(self, route): else: route.continue_() +class BaseSingleStepWorkflow(BaseUIWorkflow): + """Common interface for single-step UI workflow.""" + def __init__(self, wallet_chain_id: int, wallet_address: str, chat_message_id: str, workflow_type: str, workflow_params: Dict, runnable_step: RunnableStep) -> None: + self.chat_message_id = chat_message_id + self.workflow_type = workflow_type + self.workflow_params = workflow_params + self.runnable_step = runnable_step + self.description = runnable_step.description + + browser_storage_state = None + parsed_user_request = None + super().__init__(wallet_chain_id, wallet_address, parsed_user_request, browser_storage_state) + + def run(self) -> SingleStepResult: + try: + self.parsed_user_request = f"chat_message_id: {self.chat_message_id}, params: {self.workflow_params}" + return super().run() + except WorkflowValidationError as e: + print("UI SINGLE STEP WORKFLOW VALIDATION ERROR") + traceback.print_exc() + return SingleStepResult( + status="error", + workflow_type=self.workflow_type, + error_msg="Unexpected error. Check with support.", + user_action_type=self.runnable_step.user_action_type.name, + description=self.description, + ) + except Exception as e: + print("UI SINGLE STEP WORKFLOW EXCEPTION") + traceback.print_exc() + self.stop_listener() + return SingleStepResult( + status="error", + workflow_type=self.workflow_type, + error_msg="Unexpected error. Check with support.", + user_action_type=self.runnable_step.user_action_type.name, + description=self.description, + ) + + def _run_page(self, page, context) -> SingleStepResult: + result = self.runnable_step.function(page, context) + + # Arbitrary wait to allow for enough time for WalletConnect relay to send our client the tx data + page.wait_for_timeout(5000) + tx = self.stop_listener() + return SingleStepResult( + status=result.status, + workflow_type=self.workflow_type, + tx=tx, + error_msg=result.error_msg, + user_action_type=self.runnable_step.user_action_type.name, + description=self.description, + ) + class BaseMultiStepWorkflow(BaseUIWorkflow): """Common interface for multi-step UI workflow.""" diff --git a/ui_workflows/compound/__init__.py b/ui_workflows/compound/__init__.py new file mode 100644 index 00000000..785f91e7 --- /dev/null +++ b/ui_workflows/compound/__init__.py @@ -0,0 +1,4 @@ +from .compound_supply import CompoundSupplyWorkflow +from .compound_repay import CompoundRepayWorkflow +from .compound_borrow import CompoundBorrowWorkflow +from .compound_withdraw import CompoundWithdrawWorkflow diff --git a/ui_workflows/compound/compound_borrow.py b/ui_workflows/compound/compound_borrow.py new file mode 100644 index 00000000..e3413290 --- /dev/null +++ b/ui_workflows/compound/compound_borrow.py @@ -0,0 +1,98 @@ +import re +from logging import basicConfig, INFO +import time +import json +import uuid +import os +import requests +from typing import Any, Dict, List, Optional, Union, Literal, TypedDict, Callable +from dataclasses import dataclass, asdict + +from playwright.sync_api import TimeoutError as PlaywrightTimeoutError + +import env +from utils import TENDERLY_FORK_URL, w3 +from ..base import BaseUIWorkflow, Result, BaseSingleStepWorkflow, WorkflowStepClientPayload, StepProcessingResult, RunnableStep, tenderly_simulate_tx, setup_mock_db_objects +from database.models import ( + db_session, MultiStepWorkflow, WorkflowStep, WorkflowStepStatus, WorkflowStepUserActionType, ChatMessage, ChatSession, SystemConfig +) + +TWO_MINUTES = 120000 +TEN_SECONDS = 10000 + +class CompoundBorrowWorkflow(BaseSingleStepWorkflow): + + def __init__(self, wallet_chain_id: int, wallet_address: str, chat_message_id: str, workflow_type: str, workflow_params: Dict) -> None: + self.token = workflow_params['token'] + self.amount = workflow_params['amount'] + self.user_description = f"Borrow {self.amount} {self.token} on Compound Finance" + + step = RunnableStep("confirm_borrow", WorkflowStepUserActionType.tx, f"{self.token} confirm Borrow on Compound Finance", self.confirm_borrow) + + super().__init__(wallet_chain_id, wallet_address, chat_message_id, workflow_type, workflow_params, step) + + def _forward_rpc_node_reqs(self, route): + """Override to intercept requests to ENS API and modify response to simulate block production""" + post_body = route.request.post_data + + # Intercepting below request to modify timestamp to be 5 minutes in the future to simulate block production and allow ENS web app to not be stuck in waiting loop + if "eth_getBlockByNumber" in post_body: + curr_time_hex = hex(int(time.time()) + 300) + data = requests.post(TENDERLY_FORK_URL, data=post_body) + json_dict = data.json() + json_dict["result"]["timestamp"] = curr_time_hex + data = json_dict + res_text = json.dumps(data) + route.fulfill(body=res_text, headers={"access-control-allow-origin": "*", "access-control-allow-methods": "*", "access-control-allow-headers": "*"}) + else: + super()._forward_rpc_node_reqs(route) + + def _goto_page_and_open_walletconnect(self, page): + """Go to page and open WalletConnect modal""" + + page.goto(f"https://v2-app.compound.finance/") + + # Search for WalletConnect and open QRCode modal + page.locator("a").filter(has_text="Wallet Connect").click() + + def confirm_borrow(self, page, context) -> StepProcessingResult: + """Confirm borrow""" + # Find the token + try: + token_locators = page.get_by_text(re.compile(r".*\s{token}.*".format(token=self.token))) + except PlaywrightTimeoutError: + return StepProcessingResult( + status="error", + error_msg=f"{self.token} not available for Borrow", + ) + + # Find borrow + for i in range(4): + try: token_locators.nth(i).click() + except: continue + if page.locator("label").filter(has_text=re.compile(r"^Borrow$")).is_visible(): + page.locator("label").filter(has_text=re.compile(r"^Borrow$")).click() + break + page.locator(".close-x").click() + + # Fill the amount + try: + page.get_by_placeholder("0").fill(str(self.amount)) + except PlaywrightTimeoutError: + return StepProcessingResult( + status="error", + error_msg=f"{self.token} not available for Borrow", + ) + + # confirm borrow + try: + page.get_by_role("button", name="Borrow").click() + except PlaywrightTimeoutError: + return StepProcessingResult( + status="error", + error_msg=f"No Balance to Borrow {self.amount} {self.token}", + ) + + return StepProcessingResult( + status="success", + ) \ No newline at end of file diff --git a/ui_workflows/compound/compound_repay.py b/ui_workflows/compound/compound_repay.py new file mode 100644 index 00000000..4e6b085c --- /dev/null +++ b/ui_workflows/compound/compound_repay.py @@ -0,0 +1,114 @@ +import re +from logging import basicConfig, INFO +import time +import json +import uuid +import os +import requests +from typing import Any, Dict, List, Optional, Union, Literal, TypedDict, Callable +from dataclasses import dataclass, asdict + +from playwright.sync_api import TimeoutError as PlaywrightTimeoutError + +import env +from utils import TENDERLY_FORK_URL, w3 +from ..base import BaseUIWorkflow, MultiStepResult, BaseMultiStepWorkflow, WorkflowStepClientPayload, StepProcessingResult, RunnableStep, tenderly_simulate_tx, setup_mock_db_objects +from database.models import ( + db_session, MultiStepWorkflow, WorkflowStep, WorkflowStepStatus, WorkflowStepUserActionType, ChatMessage, ChatSession, SystemConfig +) + +TWO_MINUTES = 120000 +TEN_SECONDS = 10000 + +class CompoundRepayWorkflow(BaseMultiStepWorkflow): + + def __init__(self, wallet_chain_id: int, wallet_address: str, chat_message_id: str, workflow_type: str, workflow_params: Dict, workflow: Optional[MultiStepWorkflow] = None, curr_step_client_payload: Optional[WorkflowStepClientPayload] = None) -> None: + self.token = workflow_params['token'] + self.amount = workflow_params['amount'] + + step1 = RunnableStep("enable_repay", WorkflowStepUserActionType.tx, f"{self.token} enable Repay on Compound Finance", self.step_1_enable_repay) + step2 = RunnableStep("confirm_repay", WorkflowStepUserActionType.tx, f"{self.token} confirm Repay on Compound Finance", self.step_2_confirm_repay) + + steps = [step1, step2] + + super().__init__(wallet_chain_id, wallet_address, chat_message_id, workflow_type, workflow, workflow_params, curr_step_client_payload, steps) + + def _forward_rpc_node_reqs(self, route): + """Override to intercept requests to ENS API and modify response to simulate block production""" + post_body = route.request.post_data + + # Intercepting below request to modify timestamp to be 5 minutes in the future to simulate block production and allow ENS web app to not be stuck in waiting loop + if "eth_getBlockByNumber" in post_body: + curr_time_hex = hex(int(time.time()) + 300) + data = requests.post(TENDERLY_FORK_URL, data=post_body) + json_dict = data.json() + json_dict["result"]["timestamp"] = curr_time_hex + data = json_dict + res_text = json.dumps(data) + route.fulfill(body=res_text, headers={"access-control-allow-origin": "*", "access-control-allow-methods": "*", "access-control-allow-headers": "*"}) + else: + super()._forward_rpc_node_reqs(route) + + def _goto_page_and_open_walletconnect(self, page): + """Go to page and open WalletConnect modal""" + + page.goto(f"https://v2-app.compound.finance/") + + # Search for WalletConnect and open QRCode modal + page.locator("a").filter(has_text="Wallet Connect").click() + + def step_1_enable_repay(self, page, context) -> StepProcessingResult: + """Step 1: Enable repay""" + # Find the token + try: + token_locators = page.get_by_text(re.compile(r".*{token}.*".format(token=self.token), re.IGNORECASE)) + except PlaywrightTimeoutError: + return StepProcessingResult(status='error', error_msg=f"{self.token} not available for Repay") + + # Find Repay and enable + for i in range(4): + try: token_locators.nth(i).click() + except: continue + if page.get_by_text("Repay").is_visible(): + page.get_by_text("Repay").click() + if page.get_by_role("button", name="Enable").is_visible(): page.get_by_role("button", name="Enable").click() + # Preserve browser local storage item to allow protocol to recreate the correct state + self._preserve_browser_local_storage_item(context, 'preferences') + return StepProcessingResult(status='success') + page.locator(".close-x").click() + + return StepProcessingResult(status='error', error_msg=f"{self.token} not available for Repay") + + def step_2_confirm_repay(self, page, context) -> StepProcessingResult: + """Step 2: Confirm repay""" + # Find the token + try: + token_locators = page.get_by_text(re.compile(r".*{token}.*".format(token=self.token), re.IGNORECASE)) + except PlaywrightTimeoutError: + return StepProcessingResult(status='error', error_msg=f"{self.token} not available for Repay") + + # Find repay + for i in range(4): + try: token_locators.nth(i).click() + except: continue + if page.get_by_text("Repay").is_visible(): + page.get_by_text("Repay").click() + break + page.locator(".close-x").click() + + # Fill the amount + try: + page.get_by_placeholder("0").fill(str(self.amount)) + except PlaywrightTimeoutError: + return StepProcessingResult( + status="error", + error_msg=f"{self.token} not available for Repay", + ) + + # confirm repay + try: + page.get_by_role("button", name="Repay").click() + except PlaywrightTimeoutError: + return StepProcessingResult(status='error', error_msg=f"No Balance to Repay {self.amount} {self.token}") + + return StepProcessingResult(status='success') \ No newline at end of file diff --git a/ui_workflows/compound/compound_supply.py b/ui_workflows/compound/compound_supply.py new file mode 100644 index 00000000..b826f160 --- /dev/null +++ b/ui_workflows/compound/compound_supply.py @@ -0,0 +1,114 @@ +import re +from logging import basicConfig, INFO +import time +import json +import uuid +import os +import requests +from typing import Any, Dict, List, Optional, Union, Literal, TypedDict, Callable +from dataclasses import dataclass, asdict + +from playwright.sync_api import TimeoutError as PlaywrightTimeoutError + +import env +from utils import TENDERLY_FORK_URL, w3 +from ..base import BaseUIWorkflow, MultiStepResult, BaseMultiStepWorkflow, WorkflowStepClientPayload, StepProcessingResult, RunnableStep, tenderly_simulate_tx, setup_mock_db_objects +from database.models import ( + db_session, MultiStepWorkflow, WorkflowStep, WorkflowStepStatus, WorkflowStepUserActionType, ChatMessage, ChatSession, SystemConfig +) + +TWO_MINUTES = 120000 +TEN_SECONDS = 10000 + +class CompoundSupplyWorkflow(BaseMultiStepWorkflow): + + def __init__(self, wallet_chain_id: int, wallet_address: str, chat_message_id: str, workflow_type: str, workflow_params: Dict, workflow: Optional[MultiStepWorkflow] = None, curr_step_client_payload: Optional[WorkflowStepClientPayload] = None) -> None: + self.token = workflow_params['token'] + self.amount = workflow_params['amount'] + + step1 = RunnableStep("enable_supply", WorkflowStepUserActionType.tx, f"{self.token} enable Supply on Compound Finance", self.step_1_enable_supply) + step2 = RunnableStep("confirm_supply", WorkflowStepUserActionType.tx, f"{self.token} confirm Supply on Compound Finance", self.step_2_confirm_supply) + + steps = [step1, step2] + + super().__init__(wallet_chain_id, wallet_address, chat_message_id, workflow_type, workflow, workflow_params, curr_step_client_payload, steps) + + def _forward_rpc_node_reqs(self, route): + """Override to intercept requests to ENS API and modify response to simulate block production""" + post_body = route.request.post_data + + # Intercepting below request to modify timestamp to be 5 minutes in the future to simulate block production and allow ENS web app to not be stuck in waiting loop + if "eth_getBlockByNumber" in post_body: + curr_time_hex = hex(int(time.time()) + 300) + data = requests.post(TENDERLY_FORK_URL, data=post_body) + json_dict = data.json() + json_dict["result"]["timestamp"] = curr_time_hex + data = json_dict + res_text = json.dumps(data) + route.fulfill(body=res_text, headers={"access-control-allow-origin": "*", "access-control-allow-methods": "*", "access-control-allow-headers": "*"}) + else: + super()._forward_rpc_node_reqs(route) + + def _goto_page_and_open_walletconnect(self, page): + """Go to page and open WalletConnect modal""" + + page.goto(f"https://v2-app.compound.finance/") + + # Search for WalletConnect and open QRCode modal + page.locator("a").filter(has_text="Wallet Connect").click() + + def step_1_enable_supply(self, page, context) -> StepProcessingResult: + """Step 1: Enable supply""" + # Find the token + try: + token_locators = page.get_by_text(re.compile(r".*\s{token}.*".format(token=self.token))) + except PlaywrightTimeoutError: + return StepProcessingResult(status='error', error_msg=f"{self.token} not available for Supply") + + # Find supply and enable + for i in range(4): + try: token_locators.nth(i).click() + except: continue + if page.locator("label").filter(has_text=re.compile(r"^Supply$")).is_visible(): + page.locator("label").filter(has_text=re.compile(r"^Supply$")).click() + if page.get_by_role("button", name="Enable").is_visible(): page.get_by_role("button", name="Enable").click() + # Preserve browser local storage item to allow protocol to recreate the correct state + self._preserve_browser_local_storage_item(context, 'preferences') + return StepProcessingResult(status='success') + page.locator(".close-x").click() + + return StepProcessingResult(status='error', error_msg=f"{self.token} not available for Supply") + + def step_2_confirm_supply(self, page, context) -> StepProcessingResult: + """Step 2: Confirm supply""" + # Find the token + try: + token_locators = page.get_by_text(re.compile(r".*\s{token}.*".format(token=self.token))) + except PlaywrightTimeoutError: + return StepProcessingResult(status='error', error_msg=f"{self.token} not available for Supply") + + # Find supply + for i in range(4): + try: token_locators.nth(i).click() + except: continue + if page.locator("label").filter(has_text=re.compile(r"^Supply$")).is_visible(): + page.locator("label").filter(has_text=re.compile(r"^Supply$")).click() + break + page.locator(".close-x").click() + + # Fill the amount + # try: + page.get_by_placeholder("0").fill(str(self.amount)) + # except PlaywrightTimeoutError: + # return StepProcessingResult( + # status="error", + # error_msg=f"{self.token} not available for Supply", + # ) + + # confirm supply + try: + page.get_by_role("button", name="Supply").click() + except PlaywrightTimeoutError: + return StepProcessingResult(status='error', error_msg=f"No Balance to Supply {self.amount} {self.token}") + + return StepProcessingResult(status='success') \ No newline at end of file diff --git a/ui_workflows/compound/compound_withdraw.py b/ui_workflows/compound/compound_withdraw.py new file mode 100644 index 00000000..8c52a77b --- /dev/null +++ b/ui_workflows/compound/compound_withdraw.py @@ -0,0 +1,97 @@ +import re +from logging import basicConfig, INFO +import time +import json +import uuid +import os +import requests +from typing import Any, Dict, List, Optional, Union, Literal, TypedDict, Callable +from dataclasses import dataclass, asdict + +from playwright.sync_api import TimeoutError as PlaywrightTimeoutError + +import env +from utils import TENDERLY_FORK_URL, w3 +from ..base import BaseUIWorkflow, Result, BaseSingleStepWorkflow, WorkflowStepClientPayload, StepProcessingResult, RunnableStep, tenderly_simulate_tx, setup_mock_db_objects +from database.models import ( + db_session, MultiStepWorkflow, WorkflowStep, WorkflowStepStatus, WorkflowStepUserActionType, ChatMessage, ChatSession, SystemConfig +) + +TWO_MINUTES = 120000 +TEN_SECONDS = 10000 + +class CompoundWithdrawWorkflow(BaseSingleStepWorkflow): + + def __init__(self, wallet_chain_id: int, wallet_address: str, chat_message_id: str, workflow_type: str, workflow_params: Dict) -> None: + self.token = workflow_params['token'] + self.amount = workflow_params['amount'] + + step = RunnableStep("confirm_withdraw", WorkflowStepUserActionType.tx, f"{self.amount} {self.token} confirm Withdraw on Compound Finance", self.confirm_withdraw) + + super().__init__(wallet_chain_id, wallet_address, chat_message_id, workflow_type, workflow_params, step) + + def _forward_rpc_node_reqs(self, route): + """Override to intercept requests to ENS API and modify response to simulate block production""" + post_body = route.request.post_data + + # Intercepting below request to modify timestamp to be 5 minutes in the future to simulate block production and allow ENS web app to not be stuck in waiting loop + if "eth_getBlockByNumber" in post_body: + curr_time_hex = hex(int(time.time()) + 300) + data = requests.post(TENDERLY_FORK_URL, data=post_body) + json_dict = data.json() + json_dict["result"]["timestamp"] = curr_time_hex + data = json_dict + res_text = json.dumps(data) + route.fulfill(body=res_text, headers={"access-control-allow-origin": "*", "access-control-allow-methods": "*", "access-control-allow-headers": "*"}) + else: + super()._forward_rpc_node_reqs(route) + + def _goto_page_and_open_walletconnect(self, page): + """Go to page and open WalletConnect modal""" + + page.goto(f"https://v2-app.compound.finance/") + + # Search for WalletConnect and open QRCode modal + page.locator("a").filter(has_text="Wallet Connect").click() + + def confirm_withdraw(self, page, context) -> StepProcessingResult: + """Confirm withdraw""" + # Find the token + try: + token_locators = page.get_by_text(re.compile(r".*\s{token}.*".format(token=self.token))) + except PlaywrightTimeoutError: + return StepProcessingResult( + status="error", + error_msg=f"{self.token} not available for Withdraw", + ) + + # Find withdraw + for i in range(4): + try: token_locators.nth(i).click() + except: continue + if page.get_by_text("Withdraw").is_visible(): + page.get_by_text("Withdraw").click() + break + page.locator(".close-x").click() + + # Fill the amount + try: + page.get_by_placeholder("0").fill(str(self.amount)) + except PlaywrightTimeoutError: + return StepProcessingResult( + status="error", + error_msg=f"{self.token} not available for Withdraw", + ) + + # confirm withdraw + try: + page.get_by_role("button", name="Withdraw").click() + except PlaywrightTimeoutError: + return StepProcessingResult( + status="error", + error_msg=f"No Balance to Withdraw {self.amount} {self.token}", + ) + + return StepProcessingResult( + status="success", + ) \ No newline at end of file diff --git a/ui_workflows/compound/tests/test_compound_borrow.py b/ui_workflows/compound/tests/test_compound_borrow.py new file mode 100644 index 00000000..783e8381 --- /dev/null +++ b/ui_workflows/compound/tests/test_compound_borrow.py @@ -0,0 +1,31 @@ +import os +from ui_workflows.compound import CompoundBorrowWorkflow +from ui_workflows.base import MultiStepResult, setup_mock_db_objects, tenderly_simulate_tx +from database.models import ( + db_session, MultiStepWorkflow, WorkflowStep, WorkflowStepStatus, WorkflowStepUserActionType, ChatMessage, ChatSession, SystemConfig +) + +# Invoke this with python3 -m ui_workflows.compound.tests.test_compound_borrow +if __name__ == "__main__": + tenderly_api_access_key = os.environ.get("TENDERLY_API_ACCESS_KEY", None) + token_to_borrow = "TUSD" + wallet_address = "0xd8dA6BF26964aF9D7eEd9e03E53415D37aA96045" + wallet_chain_id = 1 # Tenderly Mainnet Fork + workflow_type = "borrow-token-on-compound" + worfklow_params = { + "token": token_to_borrow, + "amount": 0.0001, + } + mock_db_objects = setup_mock_db_objects() + mock_chat_message = mock_db_objects['mock_chat_message'] + mock_message_id = mock_chat_message.id + + print(f"Confirm Borrow of {token_to_borrow}") + + singleStepResult: SingleStepResult = CompoundBorrowWorkflow(wallet_chain_id, wallet_address, mock_message_id, workflow_type, worfklow_params).run() + + if singleStepResult.status == "success": + tenderly_simulate_tx(wallet_address, singleStepResult.tx) + print(f"Successfully Borrowed {worfklow_params['amount']} {worfklow_params['token']}") + else: + print(singleStepResult.error_msg) diff --git a/ui_workflows/compound/tests/test_compound_repay.py b/ui_workflows/compound/tests/test_compound_repay.py new file mode 100644 index 00000000..59ba22a1 --- /dev/null +++ b/ui_workflows/compound/tests/test_compound_repay.py @@ -0,0 +1,62 @@ +import os +from ui_workflows.compound import CompoundRepayWorkflow +from ui_workflows.base import MultiStepResult, setup_mock_db_objects, tenderly_simulate_tx +from database.models import ( + db_session, MultiStepWorkflow, WorkflowStep, WorkflowStepStatus, WorkflowStepUserActionType, ChatMessage, ChatSession, SystemConfig +) + +# Invoke this with python3 -m ui_workflows.compound.tests.test_compound_repay +if __name__ == "__main__": + tenderly_api_access_key = os.environ.get("TENDERLY_API_ACCESS_KEY", None) + token_to_repay = "DAI" + amount_to_repay = 0.0001 + wallet_address = "0xd8dA6BF26964aF9D7eEd9e03E53415D37aA96045" + wallet_chain_id = 1 # Tenderly Mainnet Fork + workflow_type = "repay-token-on-compound" + worfklow_params = { + "token": token_to_repay, + "amount": amount_to_repay, + } + mock_db_objects = setup_mock_db_objects() + mock_chat_message = mock_db_objects['mock_chat_message'] + mock_message_id = mock_chat_message.id + + print(f"Step 1: Enable Repay of {token_to_repay}") + + multiStepResult: MultiStepResult = CompoundRepayWorkflow(wallet_chain_id, wallet_address, mock_message_id, workflow_type, worfklow_params, None, None).run() + + # if already enabled no transaction so tx = None + if multiStepResult.tx!=None: tenderly_simulate_tx(wallet_address, multiStepResult.tx) + + print(f"Step 2: Confirm Repay of {token_to_repay}") + + workflow_id = multiStepResult.workflow_id + curr_step_client_payload = { + "id": multiStepResult.step_id, + "type": multiStepResult.step_type, + "status": multiStepResult.status, + "statusMessage": "TX successfully sent", + "userActionData": "Sample TX HASH" + } + + workflow = MultiStepWorkflow.query.filter(MultiStepWorkflow.id == workflow_id).first() + + multiStepResult: MultiStepResult = CompoundRepayWorkflow(wallet_chain_id, wallet_address, mock_message_id, workflow_type, worfklow_params, workflow, curr_step_client_payload).run() + + tenderly_simulate_tx(wallet_address, multiStepResult.tx) + + print("Final checks") + + curr_step_client_payload = { + "id": multiStepResult.step_id, + "type": multiStepResult.step_type, + "status": multiStepResult.status, + "statusMessage": "TX successfully sent", + "userActionData": "Sample TX HASH" + } + + multiStepResult: MultiStepResult = CompoundRepayWorkflow(wallet_chain_id, wallet_address, mock_message_id, workflow_type, worfklow_params, workflow, curr_step_client_payload).run() + + print(multiStepResult) + + print(f"Successfully Repaid {worfklow_params['amount']} {worfklow_params['token']}") diff --git a/ui_workflows/compound/tests/test_compound_supply.py b/ui_workflows/compound/tests/test_compound_supply.py new file mode 100644 index 00000000..0726f039 --- /dev/null +++ b/ui_workflows/compound/tests/test_compound_supply.py @@ -0,0 +1,62 @@ +import os +from ui_workflows.compound import CompoundSupplyWorkflow +from ui_workflows.base import MultiStepResult, setup_mock_db_objects, tenderly_simulate_tx +from database.models import ( + db_session, MultiStepWorkflow, WorkflowStep, WorkflowStepStatus, WorkflowStepUserActionType, ChatMessage, ChatSession, SystemConfig +) + +# Invoke this with python3 -m ui_workflows.compound.tests.test_compound_supply +if __name__ == "__main__": + tenderly_api_access_key = os.environ.get("TENDERLY_API_ACCESS_KEY", None) + token_to_supply = "LINK" + amount_to_supply = 0.0001 + wallet_address = "0xd8dA6BF26964aF9D7eEd9e03E53415D37aA96045" + wallet_chain_id = 1 # Tenderly Mainnet Fork + workflow_type = "supply-token-on-compound" + worfklow_params = { + "token": token_to_supply, + "amount": amount_to_supply, + } + mock_db_objects = setup_mock_db_objects() + mock_chat_message = mock_db_objects['mock_chat_message'] + mock_message_id = mock_chat_message.id + + print(f"Step 1: Enable Supply of {token_to_supply}") + + multiStepResult: MultiStepResult = CompoundSupplyWorkflow(wallet_chain_id, wallet_address, mock_message_id, workflow_type, worfklow_params, None, None).run() + + # if already enabled no transaction so tx = None + if multiStepResult.tx!=None: tenderly_simulate_tx(wallet_address, multiStepResult.tx) + + print(f"Step 2: Confirm Supply of {token_to_supply}") + + workflow_id = multiStepResult.workflow_id + curr_step_client_payload = { + "id": multiStepResult.step_id, + "type": multiStepResult.step_type, + "status": multiStepResult.status, + "statusMessage": "TX successfully sent", + "userActionData": "Sample TX HASH" + } + + workflow = MultiStepWorkflow.query.filter(MultiStepWorkflow.id == workflow_id).first() + + multiStepResult: MultiStepResult = CompoundSupplyWorkflow(wallet_chain_id, wallet_address, mock_message_id, workflow_type, worfklow_params, workflow, curr_step_client_payload).run() + + tenderly_simulate_tx(wallet_address, multiStepResult.tx) + + print("Final checks") + + curr_step_client_payload = { + "id": multiStepResult.step_id, + "type": multiStepResult.step_type, + "status": multiStepResult.status, + "statusMessage": "TX successfully sent", + "userActionData": "Sample TX HASH" + } + + multiStepResult: MultiStepResult = CompoundSupplyWorkflow(wallet_chain_id, wallet_address, mock_message_id, workflow_type, worfklow_params, workflow, curr_step_client_payload).run() + + print(multiStepResult) + + print(f"Successfully Supplied {worfklow_params['amount']} {worfklow_params['token']}") diff --git a/ui_workflows/compound/tests/test_compound_withdraw.py b/ui_workflows/compound/tests/test_compound_withdraw.py new file mode 100644 index 00000000..6564a6e0 --- /dev/null +++ b/ui_workflows/compound/tests/test_compound_withdraw.py @@ -0,0 +1,31 @@ +import os +from ui_workflows.compound import CompoundWithdrawWorkflow +from ui_workflows.base import MultiStepResult, setup_mock_db_objects, tenderly_simulate_tx +from database.models import ( + db_session, MultiStepWorkflow, WorkflowStep, WorkflowStepStatus, WorkflowStepUserActionType, ChatMessage, ChatSession, SystemConfig +) + +# Invoke this with python3 -m ui_workflows.compound.tests.test_compound_withdraw +if __name__ == "__main__": + tenderly_api_access_key = os.environ.get("TENDERLY_API_ACCESS_KEY", None) + token_to_withdraw = "COMP" + wallet_address = "0xd8dA6BF26964aF9D7eEd9e03E53415D37aA96045" + wallet_chain_id = 1 # Tenderly Mainnet Fork + workflow_type = "withdraw-token-on-compound" + worfklow_params = { + "token": token_to_withdraw, + "amount": 0.0001, + } + mock_db_objects = setup_mock_db_objects() + mock_chat_message = mock_db_objects['mock_chat_message'] + mock_message_id = mock_chat_message.id + + print(f"Confirm Withdraw of {token_to_withdraw}") + + singleStepResult: SingleStepResult = CompoundWithdrawWorkflow(wallet_chain_id, wallet_address, mock_message_id, workflow_type, worfklow_params).run() + + if singleStepResult.status == "success": + tenderly_simulate_tx(wallet_address, singleStepResult.tx) + print(f"Successfully Withdrawn {worfklow_params['amount']} {worfklow_params['token']}") + else: + print(singleStepResult.error_msg) diff --git a/utils.py b/utils.py index c7bb17c3..c0b85557 100644 --- a/utils.py +++ b/utils.py @@ -37,7 +37,7 @@ def set_api_key(): w3 = Web3(Web3.HTTPProvider(TENDERLY_FORK_URL)) tokenizer = tiktoken.encoding_for_model("text-davinci-003") -ns = ENS.from_web3(w3) +ns = ENS.fromWeb3(w3) def get_token_len(s: str) -> int: