Add custom exceptions for CMAB client errors and enhance error handling in fetch methods

Author: FarhanAnjum-opti

optimizely/cmab/cmab_client.py

@@ -17,6 +17,7 @@
 from typing import Dict, Any, Optional
 from optimizely import logger as _logging
 from optimizely.helpers.enums import Errors
+from optimizely.exceptions import CmabFetchError, CmabInvalidResponseError
 
 # CMAB_PREDICTION_ENDPOINT is the endpoint for CMAB predictions
 CMAB_PREDICTION_ENDPOINT = "https://prediction.cmab.optimizely.com/predict/%s"
@@ -71,20 +72,23 @@ def fetch_decision(
         rule_id: str,
         user_id: str,
         attributes: Dict[str, Any],
-        cmab_uuid: str
-    ) -> Optional[str]:
+        cmab_uuid: str,
+        timeout: Optional[float] = None
+    ) -> str:
         """Fetch a decision from the CMAB prediction service.
 
         Args:
             rule_id (str): The rule ID for the experiment.
             user_id (str): The user ID for the request.
             attributes (Dict[str, Any]): User attributes for the request.
             cmab_uuid (str): Unique identifier for the CMAB request.
+            timeout (float): Maximum wait time for request to respond in seconds.
 
         Returns:
-            Optional[str]: The variation ID if successful, None otherwise.
+            str: The variation ID.
         """
         url = CMAB_PREDICTION_ENDPOINT % rule_id
+        timeout = timeout or MAX_WAIT_TIME
         cmab_attributes = [
             {"id": key, "value": value, "type": "custom_attribute"}
             for key, value in attributes.items()
@@ -101,45 +105,50 @@ def fetch_decision(
 
         try:
             if self.retry_config:
-                variation_id = self._do_fetch_with_retry(url, request_body, self.retry_config)
+                variation_id = self._do_fetch_with_retry(url, request_body, self.retry_config, timeout)
             else:
-                variation_id = self._do_fetch(url, request_body)
+                variation_id = self._do_fetch(url, request_body, timeout)
             return variation_id
 
         except requests.RequestException as e:
-            self.logger.error(Errors.CMAB_FETCH_FAILED.format(str(e)))
-            return None
+            error_message = Errors.CMAB_FETCH_FAILED.format(str(e))
+            self.logger.error(error_message)
+            raise CmabFetchError(error_message)
 
-    def _do_fetch(self, url: str, request_body: Dict[str, Any]) -> Optional[str]:
+    def _do_fetch(self, url: str, request_body: Dict[str, Any], timeout: float) -> str:
         """Perform a single fetch request to the CMAB prediction service.
 
         Args:
             url (str): The endpoint URL.
             request_body (Dict[str, Any]): The request payload.
-
+            timeout (float): Maximum wait time for request to respond in seconds.
         Returns:
-            Optional[str]: The variation ID if successful, None otherwise.
+            str: The variation ID
         """
         headers = {'Content-Type': 'application/json'}
         try:
-            response = self.http_client.post(url, data=json.dumps(request_body), headers=headers, timeout=MAX_WAIT_TIME)
+            response = self.http_client.post(url, data=json.dumps(request_body), headers=headers, timeout=timeout)
         except requests.exceptions.RequestException as e:
-            self.logger.exception(Errors.CMAB_FETCH_FAILED.format(str(e)))
-            return None
+            error_message = Errors.CMAB_FETCH_FAILED.format(str(e))
+            self.logger.error(error_message)
+            raise CmabFetchError(error_message)
 
         if not 200 <= response.status_code < 300:
-            self.logger.exception(Errors.CMAB_FETCH_FAILED.format(str(response.status_code)))
-            return None
+            error_message = Errors.CMAB_FETCH_FAILED.format(str(response.status_code))
+            self.logger.error(error_message)
+            raise CmabFetchError(error_message)
 
         try:
             body = response.json()
         except json.JSONDecodeError:
-            self.logger.exception(Errors.INVALID_CMAB_FETCH_RESPONSE)
-            return None
+            error_message = Errors.INVALID_CMAB_FETCH_RESPONSE
+            self.logger.error(error_message)
+            raise CmabInvalidResponseError(error_message)
 
         if not self.validate_response(body):
-            self.logger.exception(Errors.INVALID_CMAB_FETCH_RESPONSE)
-            return None
+            error_message = Errors.INVALID_CMAB_FETCH_RESPONSE
+            self.logger.error(error_message)
+            raise CmabInvalidResponseError(error_message)
 
         return str(body['predictions'][0]['variation_id'])
 
@@ -165,27 +174,31 @@ def _do_fetch_with_retry(
         self,
         url: str,
         request_body: Dict[str, Any],
-        retry_config: CmabRetryConfig
-    ) -> Optional[str]:
+        retry_config: CmabRetryConfig,
+        timeout: float
+    ) -> str:
         """Perform a fetch request with retry logic.
 
         Args:
             url (str): The endpoint URL.
             request_body (Dict[str, Any]): The request payload.
             retry_config (CmabRetryConfig): Configuration for retry logic.
-
+            timeout (float): Maximum wait time for request to respond in seconds.
         Returns:
-            Optional[str]: The variation ID if successful, None otherwise.
+            str: The variation ID
         """
         backoff = retry_config.initial_backoff
         for attempt in range(retry_config.max_retries + 1):
-            variation_id = self._do_fetch(url, request_body)
-            if variation_id:
+            try:
+                variation_id = self._do_fetch(url, request_body, timeout)
                 return variation_id
-            if attempt < retry_config.max_retries:
-                self.logger.info(f"Retrying CMAB request (attempt: {attempt + 1}) after {backoff} seconds...")
-                time.sleep(backoff)
-                backoff = min(backoff * math.pow(retry_config.backoff_multiplier, attempt + 1),
-                              retry_config.max_backoff)
-        self.logger.error(Errors.CMAB_FETCH_FAILED.format('Exhausted all retries for CMAB request.'))
-        return None
+            except:
+                if attempt < retry_config.max_retries:
+                    self.logger.info(f"Retrying CMAB request (attempt: {attempt + 1}) after {backoff} seconds...")
+                    time.sleep(backoff)
+                    backoff = min(backoff * math.pow(retry_config.backoff_multiplier, attempt + 1),
+                                  retry_config.max_backoff)
+
+        error_message = Errors.CMAB_FETCH_FAILED.format('Exhausted all retries for CMAB request.')
+        self.logger.error(error_message)
+        raise CmabFetchError(error_message)

optimizely/exceptions.py

@@ -82,3 +82,21 @@ class OdpInvalidData(Exception):
     """ Raised when passing invalid ODP data. """
 
     pass
+
+
+class CmabError(Exception):
+    """Base exception for CMAB client errors."""
+
+    pass
+
+
+class CmabFetchError(CmabError):
+    """Exception raised when CMAB fetch fails."""
+
+    pass
+
+
+class CmabInvalidResponseError(CmabError):
+    """Exception raised when CMAB response is invalid."""
+
+    pass

tests/test_cmab_client.py

@@ -4,6 +4,7 @@
 from optimizely.cmab.cmab_client import DefaultCmabClient, CmabRetryConfig
 from requests.exceptions import RequestException
 from optimizely.helpers.enums import Errors
+from optimizely.exceptions import CmabFetchError, CmabInvalidResponseError
 
 
 class TestDefaultCmabClient_do_fetch(unittest.TestCase):
@@ -20,40 +21,52 @@ def test_do_fetch_success(self):
         }
         self.mock_http_client.post.return_value = mock_response
 
-        result = self.client._do_fetch('http://fake-url', {'some': 'data'})
+        result = self.client._do_fetch('http://fake-url', {'some': 'data'}, 1.0)
         self.assertEqual(result, 'abc123')
 
     def test_do_fetch_http_exception(self):
         self.mock_http_client.post.side_effect = RequestException('Connection error')
-        result = self.client._do_fetch('http://fake-url', {'some': 'data'})
-        self.assertIsNone(result)
-        self.mock_logger.exception.assert_called_with(Errors.CMAB_FETCH_FAILED.format('Connection error'))
+
+        with self.assertRaises(CmabFetchError) as context:
+            self.client._do_fetch('http://fake-url', {'some': 'data'}, 1.0)
+
+        self.mock_logger.error.assert_called_with(Errors.CMAB_FETCH_FAILED.format('Connection error'))
+        self.assertIn('Connection error', str(context.exception))
 
     def test_do_fetch_non_2xx_status(self):
         mock_response = MagicMock()
         mock_response.status_code = 500
         self.mock_http_client.post.return_value = mock_response
-        result = self.client._do_fetch('http://fake-url', {'some': 'data'})
-        self.assertIsNone(result)
-        self.mock_logger.exception.assert_called_with(Errors.CMAB_FETCH_FAILED.format(str(mock_response.status_code)))
+
+        with self.assertRaises(CmabFetchError) as context:
+            self.client._do_fetch('http://fake-url', {'some': 'data'}, 1.0)
+
+        self.mock_logger.error.assert_called_with(Errors.CMAB_FETCH_FAILED.format(str(mock_response.status_code)))
+        self.assertIn(str(mock_response.status_code), str(context.exception))
 
     def test_do_fetch_invalid_json(self):
         mock_response = MagicMock()
         mock_response.status_code = 200
         mock_response.json.side_effect = json.JSONDecodeError("Expecting value", "", 0)
         self.mock_http_client.post.return_value = mock_response
-        result = self.client._do_fetch('http://fake-url', {'some': 'data'})
-        self.assertIsNone(result)
-        self.mock_logger.exception.assert_called_with(Errors.INVALID_CMAB_FETCH_RESPONSE)
+
+        with self.assertRaises(CmabInvalidResponseError) as context:
+            self.client._do_fetch('http://fake-url', {'some': 'data'}, 1.0)
+
+        self.mock_logger.error.assert_called_with(Errors.INVALID_CMAB_FETCH_RESPONSE)
+        self.assertIn(Errors.INVALID_CMAB_FETCH_RESPONSE, str(context.exception))
 
     def test_do_fetch_invalid_response_structure(self):
         mock_response = MagicMock()
         mock_response.status_code = 200
         mock_response.json.return_value = {'no_predictions': []}
         self.mock_http_client.post.return_value = mock_response
-        result = self.client._do_fetch('http://fake-url', {'some': 'data'})
-        self.assertIsNone(result)
-        self.mock_logger.exception.assert_called_with(Errors.INVALID_CMAB_FETCH_RESPONSE)
+
+        with self.assertRaises(CmabInvalidResponseError) as context:
+            self.client._do_fetch('http://fake-url', {'some': 'data'}, 1.0)
+
+        self.mock_logger.error.assert_called_with(Errors.INVALID_CMAB_FETCH_RESPONSE)
+        self.assertIn(Errors.INVALID_CMAB_FETCH_RESPONSE, str(context.exception))
 
 
 class TestDefaultCmabClientWithRetry(unittest.TestCase):
@@ -76,7 +89,7 @@ def test_do_fetch_with_retry_success_on_first_try(self, _):
         }
         self.mock_http_client.post.return_value = mock_response
 
-        result = self.client._do_fetch_with_retry("http://fake-url", {}, self.retry_config)
+        result = self.client._do_fetch_with_retry("http://fake-url", {}, self.retry_config, 1.0)
         self.assertEqual(result, "abc123")
         self.assertEqual(self.mock_http_client.post.call_count, 1)
 
@@ -97,7 +110,7 @@ def test_do_fetch_with_retry_success_on_retry(self, _):
             success_response
         ]
 
-        result = self.client._do_fetch_with_retry("http://fake-url", {}, self.retry_config)
+        result = self.client._do_fetch_with_retry("http://fake-url", {}, self.retry_config, 1.0)
         self.assertEqual(result, "xyz456")
         self.assertEqual(self.mock_http_client.post.call_count, 2)
         self.mock_logger.info.assert_called_with("Retrying CMAB request (attempt: 1) after 0.01 seconds...")
@@ -109,8 +122,9 @@ def test_do_fetch_with_retry_exhausts_all_attempts(self, _):
 
         self.mock_http_client.post.return_value = failure_response
 
-        result = self.client._do_fetch_with_retry("http://fake-url", {}, self.retry_config)
-        self.assertIsNone(result)
+        with self.assertRaises(CmabFetchError):
+            self.client._do_fetch_with_retry("http://fake-url", {}, self.retry_config, 1.0)
+
         self.assertEqual(self.mock_http_client.post.call_count, 3)  # 1 original + 2 retries
         self.mock_logger.error.assert_called_with(
             Errors.CMAB_FETCH_FAILED.format("Exhausted all retries for CMAB request."))
@@ -124,7 +138,7 @@ def setUp(self):
         self.client = DefaultCmabClient(
             http_client=self.mock_http_client,
             logger=self.mock_logger,
-            retry_config=self.retry_config
+            retry_config=None
         )
         self.rule_id = 'test_rule'
         self.user_id = 'user123'
@@ -150,12 +164,6 @@ def test_fetch_decision_success_with_retry(self, mock_do_fetch_with_retry):
 
     @patch.object(DefaultCmabClient, '_do_fetch', side_effect=RequestException("Network error"))
     def test_fetch_decision_request_exception(self, mock_do_fetch):
-        result = self.client.fetch_decision(self.rule_id, self.user_id, self.attributes, self.cmab_uuid)
-        self.assertIsNone(result)
+        with self.assertRaises(CmabFetchError):
+            self.client.fetch_decision(self.rule_id, self.user_id, self.attributes, self.cmab_uuid)
         self.mock_logger.error.assert_called_with(Errors.CMAB_FETCH_FAILED.format("Network error"))
-
-    @patch.object(DefaultCmabClient, '_do_fetch', return_value=None)
-    def test_fetch_decision_invalid_response(self, mock_do_fetch):
-        result = self.client.fetch_decision(self.rule_id, self.user_id, self.attributes, self.cmab_uuid)
-        self.assertIsNone(result)
-        self.mock_logger.error.assert_called_once()