From 39758f44816c940bf9439a5044da94585b59e3ba Mon Sep 17 00:00:00 2001 From: codexbountylab Date: Thu, 2 Jul 2026 02:23:37 +0000 Subject: [PATCH] fix wallet_balance local wallet address resolution --- rustchain_mcp/server.py | 8 +++----- tests/test_wallet_tools.py | 27 ++++++++++++++++++++++++++- 2 files changed, 29 insertions(+), 6 deletions(-) diff --git a/rustchain_mcp/server.py b/rustchain_mcp/server.py index 2dfe021..a5339e5 100644 --- a/rustchain_mcp/server.py +++ b/rustchain_mcp/server.py @@ -216,11 +216,9 @@ def wallet_balance(wallet_id: str) -> dict: """ # First check if wallet exists in local keystore wallet = rustchain_crypto.load_wallet(wallet_id) - if wallet is None: - # Try querying by address directly - pass - - r = get_client().get(f"{RUSTCHAIN_NODE}/balance/{wallet_id}") + address = wallet["address"] if wallet else wallet_id + + r = get_client().get(f"{RUSTCHAIN_NODE}/balance/{address}") r.raise_for_status() return r.json() diff --git a/tests/test_wallet_tools.py b/tests/test_wallet_tools.py index 45aac19..d7239d2 100644 --- a/tests/test_wallet_tools.py +++ b/tests/test_wallet_tools.py @@ -636,7 +636,7 @@ def test_mcp_wallet_balance_with_mock(self, temp_keystore): """Test wallet_balance MCP tool with mocked HTTP response.""" from rustchain_mcp.server import wallet_create, wallet_balance - wallet_create(agent_name="balance-agent", password="") + created = wallet_create(agent_name="balance-agent", password="") mock_response = mock.Mock() mock_response.json.return_value = {"balance": 42.0, "wallet_id": "balance-agent"} @@ -651,6 +651,31 @@ def test_mcp_wallet_balance_with_mock(self, temp_keystore): assert isinstance(result, dict) assert result.get("balance") == 42.0 + mock_client.get.assert_called_once_with( + f"https://50.28.86.131/balance/{created['address']}" + ) + + def test_mcp_wallet_balance_accepts_direct_address(self, temp_keystore): + """Test wallet_balance still accepts a direct RTC address.""" + from rustchain_mcp.server import wallet_balance + + mock_response = mock.Mock() + mock_response.json.return_value = { + "balance": 12.5, + "wallet_id": "RTCdirect123", + } + mock_response.raise_for_status = mock.Mock() + + with mock.patch("rustchain_mcp.server.get_client") as mock_client_fn: + mock_client = mock.Mock() + mock_client.get.return_value = mock_response + mock_client_fn.return_value = mock_client + + result = wallet_balance(wallet_id="RTCdirect123") + + assert isinstance(result, dict) + assert result.get("balance") == 12.5 + mock_client.get.assert_called_once_with("https://50.28.86.131/balance/RTCdirect123") def test_mcp_wallet_history_with_mock(self, temp_keystore): """Test wallet_history MCP tool with mocked HTTP response."""