wayfinder-paths 0.1.13__py3-none-any.whl → 0.1.15__py3-none-any.whl

wayfinder_paths/strategies/stablecoin_yield_strategy/strategy.py

@@ -157,9 +157,8 @@ class StablecoinYieldStrategy(Strategy):
         main_wallet: dict[str, Any] | None = None,
         strategy_wallet: dict[str, Any] | None = None,
         web3_service=None,
-        api_key: str | None = None,
     ):
-        super().__init__(api_key=api_key)
+        super().__init__()
         merged_config: dict[str, Any] = dict(config or {})
         if main_wallet is not None:
             merged_config["main_wallet"] = main_wallet
@@ -271,7 +270,7 @@ class StablecoinYieldStrategy(Strategy):
         for token_id in self.tracked_token_ids:
             try:
                 success, balance_wei = await self.balance_adapter.get_balance(
-                    token_id=token_id,
+                    query=token_id,
                     wallet_address=strategy_address,
                 )
                 if success and balance_wei:
@@ -360,8 +359,8 @@ class StablecoinYieldStrategy(Strategy):
                     f"Gas token loaded: {gas_token_data.get('symbol', 'Unknown')}"
                 )
                 # Track gas token (but don't count it as a strategy asset)
-                if self.gas_token.get("id"):
-                    self._track_token(self.gas_token.get("id"))
+                if self.gas_token.get("token_id"):
+                    self._track_token(self.gas_token.get("token_id"))
             else:
                 logger.warning("Failed to fetch gas token info, using empty dict")
                 self.gas_token = {}
@@ -446,16 +445,28 @@ class StablecoinYieldStrategy(Strategy):
 
         llama_report = None
         if pool_ids:
-            success, llama_reports = await self.pool_adapter.get_pools_by_ids(
+            success, pool_list_response = await self.pool_adapter.get_pools_by_ids(
                 pool_ids=pool_ids
             )
-            if success:
+            if success and isinstance(pool_list_response, dict):
+                pools = pool_list_response.get("pools", [])
+                # Search for matching pool by id or constructed identifier
                 for identifier in pool_ids:
                     if not isinstance(identifier, str):
                         continue
-                    report = llama_reports.get(identifier.lower(), None)
-                    if report:
-                        llama_report = report
+                    identifier_lower = identifier.lower()
+                    for pool in pools:
+                        pool_id = pool.get("id", "").lower()
+                        pool_address = pool.get("address", "").lower()
+                        pool_chain_code = pool.get("chain_code", "").lower()
+                        constructed_id = f"{pool_chain_code}_{pool_address}"
+                        if (
+                            pool_id == identifier_lower
+                            or constructed_id == identifier_lower
+                        ):
+                            llama_report = pool
+                            break
+                    if llama_report:
                         break
 
         if self.current_pool_data is None and llama_report:
@@ -471,7 +482,14 @@ class StablecoinYieldStrategy(Strategy):
         elif llama_report and llama_report.get("apy") is not None:
             self.current_combined_apy_pct = llama_report.get("apy", 0) / 100
         elif self.current_pool_data:
-            self.current_combined_apy_pct = self.current_pool_data.get("apy", 0)
+            apy_pct = self.current_pool_data.get("combined_apy_pct")
+            if apy_pct is not None:
+                self.current_combined_apy_pct = float(apy_pct) / 100
+            else:
+                apy_val = self.current_pool_data.get("apy", 0)
+                self.current_combined_apy_pct = (
+                    float(apy_val) / 100 if apy_val is not None else 0
+                )
 
         pool_address = self.current_pool.get("address")
         chain_id = self.current_pool.get("chain", {}).get("id")
@@ -486,10 +504,10 @@ class StablecoinYieldStrategy(Strategy):
                 (
                     success,
                     current_pool_balance_raw,
-                ) = await self.balance_adapter.get_pool_balance(
-                    pool_address=pool_address,
+                ) = await self.balance_adapter.get_balance(
+                    query=pool_address,
+                    wallet_address=user_address,
                     chain_id=chain_id,
-                    user_address=user_address,
                 )
                 self.current_pool_balance = current_pool_balance_raw if success else 0
             except Exception as e:
@@ -535,7 +553,7 @@ class StablecoinYieldStrategy(Strategy):
 
         if self.usdc_token_info:
             status, raw_balance = await self.balance_adapter.get_balance(
-                token_id=self.usdc_token_info.get("token_id"),
+                query=self.usdc_token_info.get("token_id"),
                 wallet_address=self._get_strategy_wallet_address(),
             )
             if not status or not raw_balance:
@@ -578,7 +596,7 @@ class StablecoinYieldStrategy(Strategy):
             await self._refresh_tracked_balances()
 
             usdc_token_id = self.usdc_token_info.get("token_id")
-            gas_token_id = self.gas_token.get("id") if self.gas_token else None
+            gas_token_id = self.gas_token.get("token_id") if self.gas_token else None
 
             best_token_id = None
             best_balance_wei = 0
@@ -609,7 +627,7 @@ class StablecoinYieldStrategy(Strategy):
             strategy_address = self._get_strategy_wallet_address()
             try:
                 success, onchain_balance = await self.balance_adapter.get_balance(
-                    token_id=token.get("token_id"),
+                    query=token.get("token_id"),
                     wallet_address=strategy_address,
                 )
                 if success and onchain_balance:
@@ -635,7 +653,7 @@ class StablecoinYieldStrategy(Strategy):
         token_id = balance.get("token_id")
         if (
             isinstance(token_id, str)
-            and token_id.lower() == self.gas_token.get("id", "").lower()
+            and token_id.lower() == self.gas_token.get("token_id", "").lower()
         ):
             return True
 
@@ -680,7 +698,7 @@ class StablecoinYieldStrategy(Strategy):
                 "address": token_info.get("address"),
                 "chain": token_info.get("chain"),
             }
-            gas_token_id = self.gas_token.get("id")
+            gas_token_id = self.gas_token.get("token_id")
             logger.info(
                 f"Current pool set to: {token_info.get('symbol')} on {token_info.get('chain', {}).get('name')}"
             )
@@ -692,7 +710,7 @@ class StablecoinYieldStrategy(Strategy):
                     main_usdc_status,
                     main_usdc_balance,
                 ) = await self.balance_adapter.get_balance(
-                    token_id=token_info.get("token_id"),
+                    query=token_info.get("token_id"),
                     wallet_address=self._get_main_wallet_address(),
                 )
                 if main_usdc_status and main_usdc_balance is not None:
@@ -737,7 +755,7 @@ class StablecoinYieldStrategy(Strategy):
                     _,
                     main_gas_raw,
                 ) = await self.balance_adapter.get_balance(
-                    token_id=gas_token_id,
+                    query=gas_token_id,
                     wallet_address=self._get_main_wallet_address(),
                 )
                 main_gas_int = (
@@ -762,14 +780,14 @@ class StablecoinYieldStrategy(Strategy):
                     _,
                     main_gas_raw,
                 ) = await self.balance_adapter.get_balance(
-                    token_id=gas_token_id,
+                    query=gas_token_id,
                     wallet_address=self._get_main_wallet_address(),
                 )
                 (
                     _,
                     strategy_gas_raw,
                 ) = await self.balance_adapter.get_balance(
-                    token_id=gas_token_id,
+                    query=gas_token_id,
                     wallet_address=self._get_strategy_wallet_address(),
                 )
                 main_gas_int = (
@@ -895,10 +913,10 @@ class StablecoinYieldStrategy(Strategy):
             (
                 _,
                 self.current_pool_balance,
-            ) = await self.balance_adapter.get_pool_balance(
-                pool_address=self.current_pool.get("address"),
+            ) = await self.balance_adapter.get_balance(
+                query=self.current_pool.get("address"),
+                wallet_address=self._get_strategy_wallet_address(),
                 chain_id=self.current_pool.get("chain").get("id"),
-                user_address=self._get_strategy_wallet_address(),
             )
             logger.info(f"Current pool balance: {self.current_pool_balance}")
         except Exception as e:
@@ -931,8 +949,8 @@ class StablecoinYieldStrategy(Strategy):
                     )
                     if (
                         success
-                        and quotes.get("quotes")
-                        and quotes.get("quotes").get("best_quote")
+                        and isinstance(quotes, dict)
+                        and quotes.get("best_quote")
                     ):
                         logger.info("Successfully obtained swap quote")
                         break
@@ -941,7 +959,7 @@ class StablecoinYieldStrategy(Strategy):
                     if attempt == 3:  # Last attempt
                         logger.error("All quote attempts failed")
 
-            best_quote = quotes.get("quotes").get("best_quote")
+            best_quote = quotes.get("best_quote") if isinstance(quotes, dict) else None
             if not best_quote:
                 return (
                     False,
@@ -998,7 +1016,7 @@ class StablecoinYieldStrategy(Strategy):
         if self.usdc_token_info.get("token_id") in withdrawn_token_ids:
             pass
         status, raw_balance = await self.balance_adapter.get_balance(
-            token_id=self.usdc_token_info.get("token_id"),
+            query=self.usdc_token_info.get("token_id"),
             wallet_address=self._get_strategy_wallet_address(),
         )
         if not status or not raw_balance:
@@ -1025,9 +1043,9 @@ class StablecoinYieldStrategy(Strategy):
         )
         withdrawn_token_ids.add(self.usdc_token_info.get("token_id"))
 
-        if self.gas_token and self.gas_token.get("id") not in withdrawn_token_ids:
+        if self.gas_token and self.gas_token.get("token_id") not in withdrawn_token_ids:
             status, raw_gas = await self.balance_adapter.get_balance(
-                token_id=self.gas_token.get("id"),
+                query=self.gas_token.get("token_id"),
                 wallet_address=self._get_strategy_wallet_address(),
             )
             if status and raw_gas:
@@ -1039,7 +1057,7 @@ class StablecoinYieldStrategy(Strategy):
                         move_gas_status,
                         move_gas_message,
                     ) = await self.balance_adapter.move_from_strategy_wallet_to_main_wallet(
-                        self.gas_token.get("id"),
+                        self.gas_token.get("token_id"),
                         gas_amount,
                         strategy_name=self.name,
                     )
@@ -1051,7 +1069,7 @@ class StablecoinYieldStrategy(Strategy):
                                 float(gas_amount),
                             )
                         )
-                        withdrawn_token_ids.add(self.gas_token.get("id"))
+                        withdrawn_token_ids.add(self.gas_token.get("token_id"))
 
         self.DEPOSIT_USDC = 0
         self.current_pool_balance = 0
@@ -1071,6 +1089,71 @@ class StablecoinYieldStrategy(Strategy):
             f"Successfully withdrew {amount} USDC from strategy: {breakdown_msg}",
         )
 
+    async def exit(self, **kwargs) -> StatusTuple:
+        """Transfer funds from strategy wallet to main wallet."""
+        logger.info("EXIT: Transferring remaining funds to main wallet")
+
+        strategy_address = self._get_strategy_wallet_address()
+        main_address = self._get_main_wallet_address()
+
+        if strategy_address.lower() == main_address.lower():
+            return (True, "Main wallet is strategy wallet, no transfer needed")
+
+        transferred_items = []
+
+        # Transfer USDC to main wallet
+        usdc_ok, usdc_raw = await self.balance_adapter.get_balance(
+            token_id="usd-coin-base",
+            wallet_address=strategy_address,
+        )
+        if usdc_ok and usdc_raw:
+            usdc_balance = float(usdc_raw.get("balance", 0))
+            if usdc_balance > 1.0:
+                logger.info(f"Transferring {usdc_balance:.2f} USDC to main wallet")
+                (
+                    success,
+                    msg,
+                ) = await self.balance_adapter.move_from_strategy_wallet_to_main_wallet(
+                    query="usd-coin-base",
+                    amount=usdc_balance,
+                    strategy_name=self.name,
+                    skip_ledger=False,
+                )
+                if success:
+                    transferred_items.append(f"{usdc_balance:.2f} USDC")
+                else:
+                    logger.warning(f"USDC transfer failed: {msg}")
+
+        # Transfer ETH (minus reserve for tx fees) to main wallet
+        eth_ok, eth_raw = await self.balance_adapter.get_balance(
+            token_id="ethereum-base",
+            wallet_address=strategy_address,
+        )
+        if eth_ok and eth_raw:
+            eth_balance = float(eth_raw.get("balance", 0))
+            tx_fee_reserve = 0.0002
+            transferable_eth = eth_balance - tx_fee_reserve
+            if transferable_eth > 0.0001:
+                logger.info(f"Transferring {transferable_eth:.6f} ETH to main wallet")
+                (
+                    success,
+                    msg,
+                ) = await self.balance_adapter.move_from_strategy_wallet_to_main_wallet(
+                    query="ethereum-base",
+                    amount=transferable_eth,
+                    strategy_name=self.name,
+                    skip_ledger=False,
+                )
+                if success:
+                    transferred_items.append(f"{transferable_eth:.6f} ETH")
+                else:
+                    logger.warning(f"ETH transfer failed: {msg}")
+
+        if not transferred_items:
+            return (True, "No funds to transfer to main wallet")
+
+        return (True, f"Transferred to main wallet: {', '.join(transferred_items)}")
+
     async def _get_last_rotation_time(self, wallet_address: str) -> datetime | None:
         success, data = await self.ledger_adapter.get_strategy_latest_transactions(
             wallet_address=self._get_strategy_wallet_address(),
@@ -1206,7 +1289,14 @@ class StablecoinYieldStrategy(Strategy):
         else:
             self.current_pool_data = None
         if self.current_pool_data:
-            self.current_combined_apy_pct = self.current_pool_data.get("apy", 0)
+            apy_pct = self.current_pool_data.get("combined_apy_pct")
+            if apy_pct is not None:
+                self.current_combined_apy_pct = float(apy_pct) / 100
+            else:
+                apy_val = self.current_pool_data.get("apy", 0)
+                self.current_combined_apy_pct = (
+                    float(apy_val) / 100 if apy_val is not None else 0
+                )
         else:
             self.current_combined_apy_pct = (
                 target_pool_data.get("combined_apy_pct", 0) / 100
@@ -1242,10 +1332,10 @@ class StablecoinYieldStrategy(Strategy):
             (
                 _,
                 refreshed_pool_balance,
-            ) = await self.balance_adapter.get_pool_balance(
-                pool_address=pool.get("address"),
+            ) = await self.balance_adapter.get_balance(
+                query=pool.get("address"),
+                wallet_address=strategy_address,
                 chain_id=pool.get("chain").get("id"),
-                user_address=strategy_address,
             )
             self.current_pool_balance = int(refreshed_pool_balance)
         except Exception:
@@ -1259,7 +1349,7 @@ class StablecoinYieldStrategy(Strategy):
         target_token_id = target_token.get("token_id")
         target_chain = target_token.get("chain").get("code", "").lower()
         target_address = target_token.get("address", "").lower()
-        gas_token_id = self.gas_token.get("id") if self.gas_token else None
+        gas_token_id = self.gas_token.get("token_id") if self.gas_token else None
 
         # Swap all non-target, non-gas tokens to the target
         for token_id, balance_wei in list(self.tracked_balances.items()):
@@ -1278,7 +1368,7 @@ class StablecoinYieldStrategy(Strategy):
             # Get fresh balance to ensure accuracy
             try:
                 success, fresh_balance = await self.balance_adapter.get_balance(
-                    token_id=token_id,
+                    query=token_id,
                     wallet_address=self._get_strategy_wallet_address(),
                 )
                 if not success or not fresh_balance or int(fresh_balance) <= 0:
@@ -1318,7 +1408,7 @@ class StablecoinYieldStrategy(Strategy):
         # Refresh target token balance
         try:
             success, target_balance = await self.balance_adapter.get_balance(
-                token_id=target_token_id,
+                query=target_token_id,
                 wallet_address=self._get_strategy_wallet_address(),
             )
             if success and target_balance:
@@ -1335,7 +1425,7 @@ class StablecoinYieldStrategy(Strategy):
 
         required_gas = int(self.MIN_GAS * (10 ** self.gas_token.get("decimals")))
         _, current_gas = await self.balance_adapter.get_balance(
-            token_id=self.gas_token.get("id"),
+            query=self.gas_token.get("token_id"),
             wallet_address=strategy_address,
         )
         if current_gas >= required_gas:
@@ -1395,7 +1485,11 @@ class StablecoinYieldStrategy(Strategy):
         if not success:
             return 0.0
 
-        best_quote = exit_quotes.get("quotes").get("best_quote")
+        best_quote = (
+            exit_quotes.get("best_quote") if isinstance(exit_quotes, dict) else None
+        )
+        if not best_quote:
+            return None
         current_pool_usd_value = best_quote.get("output_amount")
 
         return float(
@@ -1407,7 +1501,7 @@ class StablecoinYieldStrategy(Strategy):
         # Refresh tracked balances
         await self._refresh_tracked_balances()
 
-        gas_token_id = self.gas_token.get("id") if self.gas_token else None
+        gas_token_id = self.gas_token.get("token_id") if self.gas_token else None
         results = []
 
         for token_id, balance_wei in self.tracked_balances.items():
@@ -1440,7 +1534,10 @@ class StablecoinYieldStrategy(Strategy):
         return results
 
     async def _find_best_pool(self) -> tuple[bool, dict[str, Any]]:
-        success, llama_data = await self.pool_adapter.get_pools()
+        chain_id = (self.usdc_token_info or {}).get("chain", {}).get("id")
+        if chain_id is None:
+            chain_id = 8453
+        success, llama_data = await self.pool_adapter.get_pools(chain_id=chain_id)
         if not success:
             return False, {"message": f"Failed to fetch Llama data: {llama_data}"}
 
@@ -1469,11 +1566,11 @@ class StablecoinYieldStrategy(Strategy):
                 )
                 if not target_status and candidate.get("token_id"):
                     target_status, target_pool = await self.token_adapter.get_token(
-                        token_id=candidate.get("token_id")
+                        query=candidate.get("token_id")
                     )
                 if not target_status and candidate.get("pool_id"):
                     target_status, target_pool = await self.token_adapter.get_token(
-                        token_id=candidate.get("pool_id")
+                        query=candidate.get("pool_id")
                     )
                 if not target_status:
                     continue
@@ -1515,7 +1612,8 @@ class StablecoinYieldStrategy(Strategy):
             return None
 
         try:
-            combined_apy_pct = pool_data.get("combined_apy_pct") / 100
+            apy_pct = pool_data.get("combined_apy_pct") or pool_data.get("apy") or 0
+            combined_apy_pct = float(apy_pct) / 100
             success, quotes = await self.brap_adapter.get_swap_quote(
                 from_token_address=current_token.get("address"),
                 to_token_address=token.get("address"),
@@ -1527,10 +1625,9 @@ class StablecoinYieldStrategy(Strategy):
             )
             if not success:
                 return None
-            quotes_data = quotes.get("quotes") if isinstance(quotes, dict) else None
-            if not isinstance(quotes_data, dict):
+            if not isinstance(quotes, dict):
                 return None
-            best_quote = quotes_data.get("best_quote")
+            best_quote = quotes.get("best_quote")
             if not best_quote:
                 return None
 
@@ -1582,7 +1679,7 @@ class StablecoinYieldStrategy(Strategy):
     async def _status(self) -> StatusDict:
         # Get ETH gas balance
         gas_success, gas_balance_wei = await self.balance_adapter.get_balance(
-            token_id=self.gas_token.get("id"),
+            query=self.gas_token.get("token_id"),
             wallet_address=self._get_strategy_wallet_address(),
         )
         gas_balance = (
@@ -1611,7 +1708,7 @@ class StablecoinYieldStrategy(Strategy):
 
         # Calculate total value from tracked non-gas balances
         total_value = 0.0
-        gas_token_id = self.gas_token.get("id") if self.gas_token else None
+        gas_token_id = self.gas_token.get("token_id") if self.gas_token else None
 
         for token_id, balance_wei in self.tracked_balances.items():
             if token_id == gas_token_id:
@@ -1680,7 +1777,7 @@ class StablecoinYieldStrategy(Strategy):
 
         usdc_token_id = self.usdc_token_info.get("token_id")
         usdc_decimals = self.usdc_token_info.get("decimals")
-        gas_token_id = self.gas_token.get("id") if self.gas_token else None
+        gas_token_id = self.gas_token.get("token_id") if self.gas_token else None
 
         # Check current USDC balance
         available_usdc_wei = self.tracked_balances.get(usdc_token_id, 0)
@@ -1745,7 +1842,7 @@ class StablecoinYieldStrategy(Strategy):
 
         # Refresh USDC balance after swaps
         success, usdc_wei = await self.balance_adapter.get_balance(
-            token_id=usdc_token_id,
+            query=usdc_token_id,
             wallet_address=self._get_strategy_wallet_address(),
         )
         if success and usdc_wei:
@@ -1790,7 +1887,7 @@ class StablecoinYieldStrategy(Strategy):
 
                 # Refresh USDC balance again
                 success, usdc_wei = await self.balance_adapter.get_balance(
-                    token_id=usdc_token_id,
+                    query=usdc_token_id,
                     wallet_address=self._get_strategy_wallet_address(),
                 )
                 if success and usdc_wei:

wayfinder_paths/strategies/stablecoin_yield_strategy/test_strategy.py

@@ -52,7 +52,10 @@ def strategy():
         usdc_balance_mock = AsyncMock(return_value=(True, 60000000))
         gas_balance_mock = AsyncMock(return_value=(True, 2000000000000000))
 
-        def get_balance_side_effect(token_id, wallet_address, **kwargs):
+        def get_balance_side_effect(query, wallet_address, **kwargs):
+            token_id = (
+                query if isinstance(query, str) else (query or {}).get("token_id")
+            )
             if token_id == "usd-coin-base" or token_id == "usd-coin":
                 return usdc_balance_mock.return_value
             elif token_id == "ethereum-base" or token_id == "ethereum":
@@ -449,7 +452,8 @@ async def test_sweep_wallet_uses_tracked_tokens(strategy):
     assert "token-2" in strategy.tracked_token_ids
 
     # Mock balance adapter to return fresh balances
-    async def get_balance_mock(token_id, **kwargs):
+    async def get_balance_mock(query, **kwargs):
+        token_id = query if isinstance(query, str) else (query or {}).get("token_id")
         balance = strategy.tracked_balances.get(token_id, 0)
         # Return the balance, ensuring it's an int
         return (True, int(balance) if balance else 0)
@@ -495,12 +499,11 @@ async def test_get_non_gas_balances_uses_tracked_state(strategy):
     strategy._track_token(pool_token_id, 50000000000000000000)
 
     # Mock refresh
-    strategy.balance_adapter.get_balance = AsyncMock(
-        side_effect=lambda token_id, **kwargs: (
-            True,
-            strategy.tracked_balances.get(token_id, 0),
-        )
-    )
+    def _get_balance_effect(query, **kwargs):
+        token_id = query if isinstance(query, str) else (query or {}).get("token_id")
+        return (True, strategy.tracked_balances.get(token_id, 0))
+
+    strategy.balance_adapter.get_balance = AsyncMock(side_effect=_get_balance_effect)
 
     # Get non-gas balances
     balances = await strategy._get_non_gas_balances()
@@ -519,11 +522,12 @@ async def test_partial_liquidate_uses_tracked_tokens(strategy):
     strategy._track_token("test-pool-base", 100000000000000000000)  # 100 POOL tokens
 
     # Mock balance and token adapters
+    def _get_balance_effect_partial(query, **kwargs):
+        token_id = query if isinstance(query, str) else (query or {}).get("token_id")
+        return (True, strategy.tracked_balances.get(token_id, 0))
+
     strategy.balance_adapter.get_balance = AsyncMock(
-        side_effect=lambda token_id, **kwargs: (
-            True,
-            strategy.tracked_balances.get(token_id, 0),
-        )
+        side_effect=_get_balance_effect_partial
     )
 
     strategy.token_adapter.get_token_price = AsyncMock(

wayfinder_paths/templates/adapter/README.md

@@ -77,7 +77,7 @@ async def test_get_pools():
     ):
         adapter = MyAdapter(config={})
         success, data = await adapter.get_pools(["pool-1"])
-        assert success is True
+        assert success
         assert "pools" in data
 ```
 

wayfinder_paths/templates/strategy/README.md

@@ -1,4 +1,4 @@
-# Strategy Template
+# Strategy Template
 
 This template provides the scaffolding for a new strategy. It mirrors the structure in `wayfinder_paths/strategies/...`.
 
@@ -70,7 +70,7 @@ class MyStrategy(Strategy):
             return (False, "Nothing to deposit")
 
         success, _ = await self.balance_adapter.get_balance(
-            token_id=self.config.get("token_id"),
+            query=self.config.get("token_id"),
             wallet_address=self.config.get("main_wallet", {}).get("address"),
         )
         if not success:
@@ -87,7 +87,7 @@ class MyStrategy(Strategy):
     async def _status(self) -> StatusDict:
         """Surface state back to run_strategy.py."""
         success, balance = await self.balance_adapter.get_balance(
-            token_id=self.config.get("token_id"),
+            query=self.config.get("token_id"),
             wallet_address=self.config.get("strategy_wallet", {}).get("address"),
         )
         return {

wayfinder_paths/templates/strategy/test_strategy.py

@@ -66,14 +66,15 @@ def strategy():
     #     usdc_balance_mock = AsyncMock(return_value=(True, 60000000))
     #     gas_balance_mock = AsyncMock(return_value=(True, 2000000000000000))
     #
-    #     def get_balance_side_effect(token_id, wallet_address, **kwargs):
+    #     def get_balance_side_effect(query, wallet_address, **kwargs):
+    #         token_id = query if isinstance(query, str) else (query or {}).get("token_id")
     #         if token_id == "usd-coin-base" or token_id == "usd-coin":
     #             return usdc_balance_mock.return_value
     #         elif token_id == "ethereum-base" or token_id == "ethereum":
     #             return gas_balance_mock.return_value
     #         return (True, 1000000000)
     #
-    #     s.balance_adapter.get_token_balance_for_wallet = AsyncMock(
+    #     s.balance_adapter.get_balance = AsyncMock(
     #         side_effect=get_balance_side_effect
     #     )