Refactor CMAB client: enhance docstrings for classes and methods, improve formatting, and clean up imports

Author: FarhanAnjum-opti

optimizely/cmab/cmab_client.py

@@ -1,22 +1,39 @@
+# Copyright 2025 Optimizely
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import json
 import time
 import requests
 import math
-from typing import Dict, Any, Optional, List
-from optimizely import logger  as _logging
+from typing import Dict, Any, Optional
+from optimizely import logger as _logging
 from optimizely.helpers.enums import Errors
 
 # CMAB_PREDICTION_ENDPOINT is the endpoint for CMAB predictions
-CMAB_PREDICTION_ENDPOINT = "https://prediction.cmab.optimizely.com/predict/%s" 
+CMAB_PREDICTION_ENDPOINT = "https://prediction.cmab.optimizely.com/predict/%s"
 
 # Default constants for CMAB requests
 DEFAULT_MAX_RETRIES = 3
 DEFAULT_INITIAL_BACKOFF = 0.1  # in seconds (100 ms)
-DEFAULT_MAX_BACKOFF = 10 # in seconds
+DEFAULT_MAX_BACKOFF = 10  # in seconds
 DEFAULT_BACKOFF_MULTIPLIER = 2.0
 MAX_WAIT_TIME = 10.0
 
+
 class CmabRetryConfig:
+    """Configuration for retrying CMAB requests.
+
+    Contains parameters for maximum retries, backoff intervals, and multipliers.
+    """
     def __init__(
         self,
         max_retries: int = DEFAULT_MAX_RETRIES,
@@ -29,22 +46,44 @@ def __init__(
         self.max_backoff = max_backoff
         self.backoff_multiplier = backoff_multiplier
 
+
 class DefaultCmabClient:
+    """Client for interacting with the CMAB service.
+
+    Provides methods to fetch decisions with optional retry logic.
+    """
     def __init__(self, http_client: Optional[requests.Session] = None,
                  retry_config: Optional[CmabRetryConfig] = None,
                  logger: Optional[_logging.Logger] = None):
+        """Initialize the CMAB client.
+
+        Args:
+            http_client (Optional[requests.Session]): HTTP client for making requests.
+            retry_config (Optional[CmabRetryConfig]): Configuration for retry logic.
+            logger (Optional[_logging.Logger]): Logger for logging messages.
+        """
         self.http_client = http_client or requests.Session()
         self.retry_config = retry_config
         self.logger = _logging.adapt_logger(logger or _logging.NoOpLogger())
-        
+
     def fetch_decision(
         self,
         rule_id: str,
         user_id: str,
         attributes: Dict[str, Any],
         cmab_uuid: str
     ) -> Optional[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.
 
+        Returns:
+            Optional[str]: The variation ID if successful, None otherwise.
+        """
         url = CMAB_PREDICTION_ENDPOINT % rule_id
         cmab_attributes = [
             {"id": key, "value": value, "type": "custom_attribute"}
@@ -59,7 +98,7 @@ def fetch_decision(
                 "cmabUUID": cmab_uuid,
             }]
         }
-        
+
         try:
             if self.retry_config:
                 variation_id = self._do_fetch_with_retry(url, request_body, self.retry_config)
@@ -71,7 +110,16 @@ def fetch_decision(
             self.logger.error(Errors.CMAB_FETCH_FAILED.format(str(e)))
             return None
 
-    def _do_fetch(self, url: str, request_body: str) -> Optional[str]:
+    def _do_fetch(self, url: str, request_body: Dict[str, Any]) -> Optional[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.
+
+        Returns:
+            Optional[str]: The variation ID if successful, None otherwise.
+        """
         headers = {'Content-Type': 'application/json'}
         try:
             response = self.http_client.post(url, data=json.dumps(request_body), headers=headers, timeout=MAX_WAIT_TIME)
@@ -85,7 +133,7 @@ def _do_fetch(self, url: str, request_body: str) -> Optional[str]:
 
         try:
             body = response.json()
-        except json.JSONDecodeError as e:
+        except json.JSONDecodeError:
             self.logger.exception(Errors.INVALID_CMAB_FETCH_RESPONSE)
             return None
 
@@ -94,18 +142,41 @@ def _do_fetch(self, url: str, request_body: str) -> Optional[str]:
             return None
 
         return str(body['predictions'][0]['variation_id'])
-    
-    def validate_response(self, body: dict) -> bool:
+
+    def validate_response(self, body: Dict[str, Any]) -> bool:
+        """Validate the response structure from the CMAB service.
+
+        Args:
+            body (Dict[str, Any]): The response body to validate.
+
+        Returns:
+            bool: True if the response is valid, False otherwise.
+        """
         return (
-            isinstance(body, dict)
-            and 'predictions' in body
-            and isinstance(body['predictions'], list)
-            and len(body['predictions']) > 0
-            and isinstance(body['predictions'][0], dict)
-            and "variation_id" in body["predictions"][0]
+            isinstance(body, dict) and
+            'predictions' in body and
+            isinstance(body['predictions'], list) and
+            len(body['predictions']) > 0 and
+            isinstance(body['predictions'][0], dict) and
+            "variation_id" in body["predictions"][0]
         )
 
-    def _do_fetch_with_retry(self, url: str, request_body: dict, retry_config: CmabRetryConfig) -> Optional[str]:
+    def _do_fetch_with_retry(
+        self,
+        url: str,
+        request_body: Dict[str, Any],
+        retry_config: CmabRetryConfig
+    ) -> Optional[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.
+
+        Returns:
+            Optional[str]: The variation ID if successful, None otherwise.
+        """
         backoff = retry_config.initial_backoff
         for attempt in range(retry_config.max_retries + 1):
             variation_id = self._do_fetch(url, request_body)
@@ -114,6 +185,7 @@ def _do_fetch_with_retry(self, url: str, request_body: dict, retry_config: CmabR
             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)
+                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

tests/test_cmab_client.py

@@ -1,10 +1,11 @@
 import unittest
 import json
-from unittest.mock import MagicMock, patch, call
+from unittest.mock import MagicMock, patch
 from optimizely.cmab.cmab_client import DefaultCmabClient, CmabRetryConfig
 from requests.exceptions import RequestException
 from optimizely.helpers.enums import Errors
 
+
 class TestDefaultCmabClient_do_fetch(unittest.TestCase):
     def setUp(self):
         self.mock_http_client = MagicMock()
@@ -54,6 +55,7 @@ def test_do_fetch_invalid_response_structure(self):
         self.assertIsNone(result)
         self.mock_logger.exception.assert_called_with(Errors.INVALID_CMAB_FETCH_RESPONSE)
 
+
 class TestDefaultCmabClientWithRetry(unittest.TestCase):
     def setUp(self):
         self.mock_http_client = MagicMock()
@@ -110,7 +112,9 @@ def test_do_fetch_with_retry_exhausts_all_attempts(self, _):
         result = self.client._do_fetch_with_retry("http://fake-url", {}, self.retry_config)
         self.assertIsNone(result)
         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."))
+        self.mock_logger.error.assert_called_with(
+            Errors.CMAB_FETCH_FAILED.format("Exhausted all retries for CMAB request."))
+
 
 class TestDefaultCmabClientFetchDecision(unittest.TestCase):
     def setUp(self):