fix: Improve remote evaluation fetch retry logic (#42)

Author: tyiuhc

.github/workflows/publish-to-pypi.yml

@@ -36,7 +36,7 @@ jobs:
           python-version: 3.7
 
       - name: Install dependencies
-        run: python -m pip install build setuptools wheel twine amplitude_analytics python-semantic-release==7.34.6
+        run: python -m pip install build setuptools wheel twine amplitude_analytics parameterized python-semantic-release==7.34.6
 
       - name: Run Test
         run: python -m unittest discover -s ./tests -p '*_test.py'

.github/workflows/publish-to-test-pypi.yml

@@ -27,7 +27,7 @@ jobs:
           python-version: 3.7
 
       - name: Install dependencies
-        run: python -m pip install build setuptools wheel twine amplitude_analytics
+        run: python -m pip install build setuptools wheel twine amplitude_analytics parameterized
 
       - name: Build a binary wheel and a source tarball
         run: python -m build --sdist --wheel --outdir dist/ .

.github/workflows/test-arm.yml

@@ -21,4 +21,5 @@ jobs:
             apt -y install ca-certificates
           run: |
             pip install -r requirements.txt
+            pip install -r requirements-dev.txt
             python3 -m unittest discover -s ./tests -p '*_test.py'

.github/workflows/test.yml

@@ -20,6 +20,7 @@ jobs:
 
       - name: Install requirements
         run: pip install -r requirements.txt
+          pip install -r requirements-dev.txt
 
       - name: Unit Test
         run: python -m unittest discover -s ./tests -p '*_test.py'

requirements-dev.txt

@@ -0,0 +1 @@
+parameterized~=0.9.0

src/amplitude_experiment/exception.py

@@ -0,0 +1,4 @@
+class FetchException(Exception):
+    def __init__(self, status_code, message):
+        super().__init__(message)
+        self.status_code = status_code

src/amplitude_experiment/remote/client.py

@@ -7,6 +7,7 @@
 
 from .config import RemoteEvaluationConfig
 from ..connection_pool import HTTPConnectionPool
+from ..exception import FetchException
 from ..user import User
 from ..util.deprecated import deprecated
 from ..util.variant import evaluation_variants_json_to_variants
@@ -112,7 +113,8 @@ def __fetch_internal(self, user):
             return self.__do_fetch(user)
         except Exception as e:
             self.logger.error(f"[Experiment] Fetch failed: {e}")
-            return self.__retry_fetch(user)
+            if self.__should_retry_fetch(e):
+                return self.__retry_fetch(user)
 
     def __retry_fetch(self, user):
         if self.config.fetch_retries == 0:
@@ -148,6 +150,9 @@ def __do_fetch(self, user):
             response = conn.request('POST', '/sdk/v2/vardata?v=0', body, headers)
             elapsed = '%.3f' % ((time.time() - start) * 1000)
             self.logger.debug(f"[Experiment] Fetch complete in {elapsed} ms")
+            if response.status != 200:
+                raise FetchException(response.status,
+                                     f"Fetch error response: status={response.status} {response.reason}")
             json_response = json.loads(response.read().decode("utf8"))
             variants = evaluation_variants_json_to_variants(json_response)
             self.logger.debug(f"[Experiment] Fetched variants: {json.dumps(variants, default=str)}")
@@ -189,6 +194,11 @@ def is_default_variant(variant: Variant) -> bool:
             if variant.metadata is not None and variant.metadata.get('deployed') is not None:
                 deployed = variant.metadata.get('deployed')
             return default and not deployed
-        return {key: variant for key, variant in variants.items() if not is_default_variant(variant)}
 
+        return {key: variant for key, variant in variants.items() if not is_default_variant(variant)}
 
+    @staticmethod
+    def __should_retry_fetch(err: Exception):
+        if isinstance(err, FetchException):
+            return err.status_code < 400 or err.status_code >= 500 or err.status_code == 429
+        return True

tests/remote/client_test.py

@@ -1,6 +1,10 @@
 import unittest
+from unittest import mock
+
+from parameterized import parameterized
 
 from src.amplitude_experiment import RemoteEvaluationClient, Variant, User, RemoteEvaluationConfig
+from src.amplitude_experiment.exception import FetchException
 
 API_KEY = 'client-DvWljIjiiuqLbyjqdvBaLFfEBrAvGuA3'
 SERVER_URL = 'https://api.lab.amplitude.com/sdk/vardata'
@@ -36,11 +40,33 @@ def test_fetch_async(self):
         self.client.fetch_async(user, self.callback_for_async)
 
     def test_fetch_failed_with_retry(self):
-        with RemoteEvaluationClient(API_KEY, RemoteEvaluationConfig(debug=False, fetch_retries=1, fetch_timeout_millis=1)) as client:
+        with RemoteEvaluationClient(API_KEY, RemoteEvaluationConfig(debug=False, fetch_retries=1,
+                                                                    fetch_timeout_millis=1)) as client:
             user = User(user_id='test_user')
             variants = client.fetch(user)
             self.assertEqual({}, variants)
 
+    @parameterized.expand([
+        (300, "Fetch Exception 300", True),
+        (400, "Fetch Exception 400", False),
+        (429, "Fetch Exception 429", True),
+        (500, "Fetch Exception 500", True),
+        (000, "Other Exception", True),
+    ])
+    @mock.patch("src.amplitude_experiment.remote.client.RemoteEvaluationClient._RemoteEvaluationClient__retry_fetch")
+    @mock.patch("src.amplitude_experiment.remote.client.RemoteEvaluationClient._RemoteEvaluationClient__do_fetch")
+    def test_fetch_retry_with_response(self, response_code, error_message, should_call_retry, mock_do_fetch,
+                                       mock_retry_fetch):
+        if response_code == 000:
+            mock_do_fetch.side_effect = Exception(error_message)
+        else:
+            mock_do_fetch.side_effect = FetchException(response_code, error_message)
+        instance = RemoteEvaluationClient(API_KEY, RemoteEvaluationConfig(fetch_retries=1))
+        user = User(user_id='test_user')
+        instance.fetch(user)
+        mock_do_fetch.assert_called_once_with(user)
+        self.assertEqual(should_call_retry, mock_retry_fetch.called)
+
 
 if __name__ == '__main__':
     unittest.main()