Merge pull request #128 from okxapi/dev/BROK-1214

Dev/brok 1214

Author: okxapi

okx/Account.py

@@ -325,3 +325,9 @@ def set_auto_repay(self, autoRepay=False):
     def spot_borrow_repay_history(self, ccy='', type='', after='', before='', limit=''):
         params = {'ccy': ccy, 'type': type, 'after': after, 'before': before, 'limit': limit}
         return self._request_with_params(GET, GET_BORROW_REPAY_HISTORY, params)
+
+    def set_auto_earn(self, ccy, action, earnType=None):
+        params = {'ccy': ccy, 'action': action}
+        if earnType is not None:
+            params['earnType'] = earnType
+        return self._request_with_params(POST, SET_AUTO_EARN, params)

okx/consts.py

@@ -66,6 +66,7 @@
 MANUAL_REBORROW_REPAY = '/api/v5/account/spot-manual-borrow-repay'
 SET_AUTO_REPAY='/api/v5/account/set-auto-repay'
 GET_BORROW_REPAY_HISTORY='/api/v5/account/spot-borrow-repay-history'
+SET_AUTO_EARN='/api/v5/account/set-auto-earn'
 
 # Funding
 NON_TRADABLE_ASSETS = '/api/v5/asset/non-tradable-assets'

test/AccountTest.py

@@ -146,6 +146,8 @@ def setUp(self):
     #     logger.info(f'{self.AccountAPI.set_auto_repay(autoRepay=True)}')
     # def test_spot_borrow_repay_history(self):
     #     logger.debug(self.AccountAPI.spot_borrow_repay_history(ccy="USDT",type="auto_borrow",after="1597026383085"))
+    def test_set_auto_earn(self):
+        logger.debug(self.AccountAPI.set_auto_earn(ccy="USDT", action="turn_on", earnType='0'))
 
 if __name__ == '__main__':
     unittest.main()

test/unit/okx/test_account.py

@@ -436,6 +436,122 @@ def test_zero_idxVol_is_included(self, mock_request):
         mock_request.assert_called_once_with(c.POST, c.POSITION_BUILDER, expected_params)
 
 
+class TestAccountAPISetAutoEarn(unittest.TestCase):
+    """Unit tests for the set_auto_earn method"""
+
+    def setUp(self):
+        """Set up test fixtures"""
+        self.account_api = AccountAPI(
+            api_key='test_key',
+            api_secret_key='test_secret',
+            passphrase='test_pass',
+            flag='0'
+        )
+
+    @patch.object(AccountAPI, '_request_with_params')
+    def test_set_auto_earn_with_all_params(self, mock_request):
+        """Test set_auto_earn with all parameters provided"""
+        # Arrange
+        mock_response = {'code': '0', 'msg': '', 'data': []}
+        mock_request.return_value = mock_response
+
+        # Act
+        result = self.account_api.set_auto_earn(
+            ccy='USDT',
+            action='turn_on',
+            earnType='0'
+        )
+
+        # Assert
+        expected_params = {
+            'ccy': 'USDT',
+            'action': 'turn_on',
+            'earnType': '0'
+        }
+        mock_request.assert_called_once_with(c.POST, c.SET_AUTO_EARN, expected_params)
+        self.assertEqual(result, mock_response)
+
+    @patch.object(AccountAPI, '_request_with_params')
+    def test_set_auto_earn_turn_on_action(self, mock_request):
+        """Test set_auto_earn with turn_on action"""
+        # Arrange
+        mock_response = {'code': '0', 'msg': '', 'data': []}
+        mock_request.return_value = mock_response
+
+        # Act
+        result = self.account_api.set_auto_earn(
+            ccy='BTC',
+            action='turn_on',
+            earnType='0'
+        )
+
+        # Assert
+        expected_params = {
+            'ccy': 'BTC',
+            'action': 'turn_on',
+            'earnType': '0'
+        }
+        mock_request.assert_called_once_with(c.POST, c.SET_AUTO_EARN, expected_params)
+
+    @patch.object(AccountAPI, '_request_with_params')
+    def test_set_auto_earn_turn_off_action(self, mock_request):
+        """Test set_auto_earn with turn_off action"""
+        # Arrange
+        mock_response = {'code': '0', 'msg': '', 'data': []}
+        mock_request.return_value = mock_response
+
+        # Act
+        result = self.account_api.set_auto_earn(
+            ccy='ETH',
+            action='turn_off',
+            earnType='0'
+        )
+
+        # Assert
+        expected_params = {
+            'ccy': 'ETH',
+            'action': 'turn_off',
+            'earnType': '0'
+        }
+        mock_request.assert_called_once_with(c.POST, c.SET_AUTO_EARN, expected_params)
+
+    @patch.object(AccountAPI, '_request_with_params')
+    def test_set_auto_earn_with_required_params_only(self, mock_request):
+        """Test set_auto_earn with required parameters only (ccy, action)"""
+        # Arrange
+        mock_response = {'code': '0', 'msg': '', 'data': []}
+        mock_request.return_value = mock_response
+
+        # Act
+        result = self.account_api.set_auto_earn(ccy='USDT', action='turn_on')
+
+        # Assert
+        expected_params = {
+            'ccy': 'USDT',
+            'action': 'turn_on'
+        }
+        mock_request.assert_called_once_with(c.POST, c.SET_AUTO_EARN, expected_params)
+
+    @patch.object(AccountAPI, '_request_with_params')
+    def test_set_auto_earn_different_currencies(self, mock_request):
+        """Test set_auto_earn with different currencies"""
+        mock_response = {'code': '0', 'msg': '', 'data': []}
+        mock_request.return_value = mock_response
+
+        currencies = ['USDT', 'BTC', 'ETH', 'USDC']
+
+        for ccy in currencies:
+            mock_request.reset_mock()
+            result = self.account_api.set_auto_earn(
+                ccy=ccy,
+                action='turn_on',
+                earnType='0'
+            )
+
+            call_args = mock_request.call_args[0][2]
+            self.assertEqual(call_args['ccy'], ccy)
+
+
 if __name__ == '__main__':
     unittest.main()