diff --git a/tests/functional_tests/test_framework/wallet.py b/tests/functional_tests/test_framework/wallet.py
index 7b19c4d55..e9f13667b 100644
--- a/tests/functional_tests/test_framework/wallet.py
+++ b/tests/functional_tests/test_framework/wallet.py
@@ -125,9 +125,14 @@ class Wallet(object):
         }
         return self.rpc.send_json_rpc_request(create_wallet)
 
-    def get_balance(self):
+    def get_balance(self, account_index = 0, address_indices = [], all_accounts = False):
         get_balance = {
             'method': 'get_balance',
+            'params': {
+                'account_index': account_index,
+                'address_indices': address_indices,
+                'all_accounts': all_accounts,
+            },
             'jsonrpc': '2.0', 
             'id': '0'
         }