fix: use ServerZone enum for remote evaluation config

Author: tyiuhc

src/amplitude_experiment/__init__.py

@@ -12,6 +12,6 @@
 from .cookie import AmplitudeCookie
 from .local.client import LocalEvaluationClient
 from .local.config import LocalEvaluationConfig
-from .local.config import ServerZone
+from .server_zone import ServerZone
 from .assignment import AssignmentConfig
 from .cohort.cohort_sync_config import CohortSyncConfig

src/amplitude_experiment/local/config.py

@@ -1,7 +1,6 @@
-from enum import Enum
-
 from ..assignment import AssignmentConfig
 from ..cohort.cohort_sync_config import CohortSyncConfig, DEFAULT_COHORT_SYNC_URL, EU_COHORT_SYNC_URL
+from ..server_zone import ServerZone
 
 DEFAULT_SERVER_URL = 'https://api.lab.amplitude.com'
 EU_SERVER_URL = 'https://flag.lab.eu.amplitude.com'
@@ -10,11 +9,6 @@
 EU_STREAM_SERVER_URL = 'https://stream.lab.eu.amplitude.com'
 
 
-class ServerZone(Enum):
-    US = "US"
-    EU = "EU"
-
-
 class LocalEvaluationConfig:
     """Experiment Local Client Configuration"""
 

src/amplitude_experiment/remote/config.py

@@ -1,3 +1,5 @@
+from ..server_zone import ServerZone
+
 DEFAULT_SERVER_URL = 'https://api.lab.amplitude.com'
 EU_SERVER_URL = 'https://api.lab.eu.amplitude.com'
 
@@ -12,7 +14,7 @@ def __init__(self, debug=False,
                  fetch_retry_backoff_max_millis=10000,
                  fetch_retry_backoff_scalar=1.5,
                  fetch_retry_timeout_millis=10000,
-                 server_zone='us'):
+                 server_zone: ServerZone = ServerZone.US):
         """
         Initialize a config
             Parameters:
@@ -27,7 +29,7 @@ def __init__(self, debug=False,
                   greater than the max, the max is used for all subsequent retries.
                 fetch_retry_backoff_scalar (float): Scales the minimum backoff exponentially.
                 fetch_retry_timeout_millis (int): The request timeout for retrying fetch requests.
-                server_zone (str): Select the Amplitude data center to get flags and variants from, `us` or `eu`.
+                server_zone (str): Select the Amplitude data center to get flags and variants from, US or EU.
 
             Returns:
                 The config object
@@ -41,6 +43,6 @@ def __init__(self, debug=False,
         self.fetch_retry_backoff_scalar = fetch_retry_backoff_scalar
         self.fetch_retry_timeout_millis = fetch_retry_timeout_millis
         self.server_zone = server_zone
-        if server_url == DEFAULT_SERVER_URL and server_zone == 'eu':
+        if server_url == DEFAULT_SERVER_URL and server_zone == ServerZone.EU:
             self.server_url = EU_SERVER_URL
 

src/amplitude_experiment/server_zone.py

@@ -0,0 +1,5 @@
+from enum import Enum
+
+class ServerZone(Enum):
+    US = "US"
+    EU = "EU"