Add account_id parameter to get_balances method for single account filtering

luno_python/client.py

@@ -156,7 +156,7 @@ def delete_beneficiary(self, id):
         }
         return self.do('DELETE', '/api/1/beneficiaries/{id}', req=req, auth=True)
 
-    def get_balances(self, assets=None):
+    def get_balances(self, assets=None, account_id=None):
         """Makes a call to GET /api/1/balance.
 
         The list of all Accounts and their respective balances for the requesting user.
@@ -168,11 +168,26 @@ def get_balances(self, assets=None):
                        pass the parameter multiple times,
                        e.g. `assets=XBT&assets=ETH`.
         :type assets: list
+        :param account_id: Only return balance for the account with this ID. If provided,
+                          returns a single account object instead of the full response.
+        :type account_id: str
         """
         req = {
             'assets': assets,
         }
-        return self.do('GET', '/api/1/balance', req=req, auth=True)
+        response = self.do('GET', '/api/1/balance', req=req, auth=True)
+
+        # If account_id is specified, filter to return only that account
+        if account_id is not None:
+            if 'balance' in response:
+                for account in response['balance']:
+                    if str(account.get('account_id')) == str(account_id):
+                        return account
+            # If account_id not found, return None
+            return None
+
+        # Return full response if no account_id specified (backward compatibility)
+        return response
 
     def get_candles(self, duration, pair, since):
         """Makes a call to GET /api/exchange/1/candles.

tests/test_client.py

@@ -48,3 +48,174 @@ def test_client_do_basic():
         res = c.do('GET', '/')
     assert e.value.code == 'code'
     assert e.value.message == 'message'
+
+
+def test_get_balances_without_account_id():
+    """Test get_balances without account_id parameter (backward compatibility)"""
+    c = Client()
+    c.set_base_url('mock://test/')
+
+    adapter = requests_mock.Adapter()
+    c.session.mount('mock', adapter)
+
+    # Mock the API response
+    mock_response = {
+        'balance': [
+            {'account_id': '12345678910', 'asset': 'XBT', 'balance': '0.00', 'reserved': '0.00', 'unconfirmed': '0.00'},
+            {'account_id': '98765432100', 'asset': 'ETH', 'balance': '1.50', 'reserved': '0.10', 'unconfirmed': '0.05'},
+            {'account_id': '55555555555', 'asset': 'ZAR', 'balance': '1000.00', 'reserved': '0.00', 'unconfirmed': '0.00'}
+        ]
+    }
+
+    adapter.register_uri('GET', 'mock://test/api/1/balance', json=mock_response)
+
+    # Test without account_id - should return full response
+    result = c.get_balances()
+    assert result == mock_response
+    assert 'balance' in result
+    assert len(result['balance']) == 3
+
+
+def test_get_balances_with_valid_account_id():
+    """Test get_balances with valid account_id parameter"""
+    c = Client()
+    c.set_base_url('mock://test/')
+
+    adapter = requests_mock.Adapter()
+    c.session.mount('mock', adapter)
+
+    # Mock the API response
+    mock_response = {
+        'balance': [
+            {'account_id': '12345678910', 'asset': 'XBT', 'balance': '0.00', 'reserved': '0.00', 'unconfirmed': '0.00'},
+            {'account_id': '98765432100', 'asset': 'ETH', 'balance': '1.50', 'reserved': '0.10', 'unconfirmed': '0.05'},
+            {'account_id': '55555555555', 'asset': 'ZAR', 'balance': '1000.00', 'reserved': '0.00', 'unconfirmed': '0.00'}
+        ]
+    }
+
+    adapter.register_uri('GET', 'mock://test/api/1/balance', json=mock_response)
+
+    # Test with valid account_id - should return single account
+    result = c.get_balances(account_id='12345678910')
+    expected = {'account_id': '12345678910', 'asset': 'XBT', 'balance': '0.00', 'reserved': '0.00', 'unconfirmed': '0.00'}
+    assert result == expected
+
+    # Test with another valid account_id
+    result = c.get_balances(account_id='98765432100')
+    expected = {'account_id': '98765432100', 'asset': 'ETH', 'balance': '1.50', 'reserved': '0.10', 'unconfirmed': '0.05'}
+    assert result == expected
+
+
+def test_get_balances_with_invalid_account_id():
+    """Test get_balances with invalid account_id parameter"""
+    c = Client()
+    c.set_base_url('mock://test/')
+
+    adapter = requests_mock.Adapter()
+    c.session.mount('mock', adapter)
+
+    # Mock the API response
+    mock_response = {
+        'balance': [
+            {'account_id': '12345678910', 'asset': 'XBT', 'balance': '0.00', 'reserved': '0.00', 'unconfirmed': '0.00'},
+            {'account_id': '98765432100', 'asset': 'ETH', 'balance': '1.50', 'reserved': '0.10', 'unconfirmed': '0.05'}
+        ]
+    }
+
+    adapter.register_uri('GET', 'mock://test/api/1/balance', json=mock_response)
+
+    # Test with invalid account_id - should return None
+    result = c.get_balances(account_id='99999999999')
+    assert result is None
+
+
+def test_get_balances_with_account_id_and_assets():
+    """Test get_balances with both account_id and assets parameters"""
+    c = Client()
+    c.set_base_url('mock://test/')
+
+    adapter = requests_mock.Adapter()
+    c.session.mount('mock', adapter)
+
+    # Mock the API response
+    mock_response = {
+        'balance': [
+            {'account_id': '12345678910', 'asset': 'XBT', 'balance': '0.00', 'reserved': '0.00', 'unconfirmed': '0.00'},
+            {'account_id': '98765432100', 'asset': 'ETH', 'balance': '1.50', 'reserved': '0.10', 'unconfirmed': '0.05'}
+        ]
+    }
+
+    adapter.register_uri('GET', 'mock://test/api/1/balance', json=mock_response)
+
+    # Test with both parameters
+    result = c.get_balances(assets=['XBT'], account_id='12345678910')
+    expected = {'account_id': '12345678910', 'asset': 'XBT', 'balance': '0.00', 'reserved': '0.00', 'unconfirmed': '0.00'}
+    assert result == expected
+
+
+def test_get_balances_with_account_id_type_conversion():
+    """Test get_balances with account_id type conversion (string vs int)"""
+    c = Client()
+    c.set_base_url('mock://test/')
+
+    adapter = requests_mock.Adapter()
+    c.session.mount('mock', adapter)
+
+    # Mock the API response with integer account_id
+    mock_response = {
+        'balance': [
+            {'account_id': 12345678910, 'asset': 'XBT', 'balance': '0.00', 'reserved': '0.00', 'unconfirmed': '0.00'},
+            {'account_id': 98765432100, 'asset': 'ETH', 'balance': '1.50', 'reserved': '0.10', 'unconfirmed': '0.05'}
+        ]
+    }
+
+    adapter.register_uri('GET', 'mock://test/api/1/balance', json=mock_response)
+
+    # Test with string account_id when API returns integer - should work due to type conversion
+    result = c.get_balances(account_id='12345678910')
+    expected = {'account_id': 12345678910, 'asset': 'XBT', 'balance': '0.00', 'reserved': '0.00', 'unconfirmed': '0.00'}
+    assert result == expected
+
+
+def test_get_balances_with_empty_balance_response():
+    """Test get_balances when API returns empty balance list"""
+    c = Client()
+    c.set_base_url('mock://test/')
+
+    adapter = requests_mock.Adapter()
+    c.session.mount('mock', adapter)
+
+    # Mock empty response
+    mock_response = {'balance': []}
+
+    adapter.register_uri('GET', 'mock://test/api/1/balance', json=mock_response)
+
+    # Test with account_id on empty response
+    result = c.get_balances(account_id='12345678910')
+    assert result is None
+
+    # Test without account_id on empty response
+    result = c.get_balances()
+    assert result == mock_response
+
+
+def test_get_balances_with_malformed_response():
+    """Test get_balances when API returns malformed response"""
+    c = Client()
+    c.set_base_url('mock://test/')
+
+    adapter = requests_mock.Adapter()
+    c.session.mount('mock', adapter)
+
+    # Mock response without 'balance' key
+    mock_response = {'some_other_key': 'value'}
+
+    adapter.register_uri('GET', 'mock://test/api/1/balance', json=mock_response)
+
+    # Test with account_id on malformed response
+    result = c.get_balances(account_id='12345678910')
+    assert result is None
+
+    # Test without account_id on malformed response
+    result = c.get_balances()
+    assert result == mock_response