optimizely · pthompson127 · Jun 17, 2020 · Jun 9, 2020 · Jun 9, 2020 · Jun 9, 2020
diff --git a/optimizely/config_manager.py b/optimizely/config_manager.py
@@ -150,6 +150,8 @@ def get_config(self):
 class PollingConfigManager(StaticConfigManager):
     """ Config manager that polls for the datafile and updated ProjectConfig based on an update interval. """
 
+    DATAFILE_URL_TEMPLATE = enums.ConfigManager.DATAFILE_URL_TEMPLATE
+
     def __init__(
         self,
         sdk_key=None,
@@ -192,7 +194,7 @@ def __init__(
             skip_json_validation=skip_json_validation,
         )
         self.datafile_url = self.get_datafile_url(
-            sdk_key, url, url_template or enums.ConfigManager.DATAFILE_URL_TEMPLATE
+            sdk_key, url, url_template or self.DATAFILE_URL_TEMPLATE
         )
         self.set_update_interval(update_interval)
         self.set_blocking_timeout(blocking_timeout)
@@ -368,3 +370,46 @@ def start(self):
         """ Start the config manager and the thread to periodically fetch datafile. """
         if not self.is_running:
             self._polling_thread.start()
+
+
+class AuthDatafilePollingConfigManager(PollingConfigManager):
+    """ Config manager that polls for authenticated datafile using access token. """
+
+    DATAFILE_URL_TEMPLATE = enums.ConfigManager.AUTHENTICATED_DATAFILE_URL_TEMPLATE
+
+    def __init__(
+        self,
+        access_token,
+        *args,
+        **kwargs
+    ):
+        """ Initialize config manager. One of sdk_key or url has to be set to be able to use.
+
+        Args:
+            access_token: String to be attached to the request header to fetch the authenticated datafile.
+            *args: Refer to arguments descriptions in PollingConfigManager.
+            **kwargs: Refer to keyword arguments descriptions in PollingConfigManager.
+        """
+        self._set_access_token(access_token)
+        super(AuthDatafilePollingConfigManager, self).__init__(*args, **kwargs)
+
+    def _set_access_token(self, access_token):
+        """ Checks for valid access token input and sets it. """
+        if not access_token:
+            raise optimizely_exceptions.InvalidInputException(
+                'access_token cannot be empty or None.')
+        self.access_token = access_token
+
+    def fetch_datafile(self):
+        """ Fetch authenticated datafile and set ProjectConfig. """
+        request_headers = {}
+        request_headers[enums.HTTPHeaders.AUTHORIZATION] = \
+            enums.ConfigManager.AUTHORIZATION_HEADER_DATA_TEMPLATE.format(access_token=self.access_token)
+
+        if self.last_modified:
+            request_headers[enums.HTTPHeaders.IF_MODIFIED_SINCE] = self.last_modified
+
+        response = requests.get(
+            self.datafile_url, headers=request_headers, timeout=enums.ConfigManager.REQUEST_TIMEOUT,
+        )
+        self._handle_response(response)
diff --git a/optimizely/helpers/enums.py b/optimizely/helpers/enums.py
@@ -49,6 +49,8 @@ class AudienceEvaluationLogs(object):
 
 
 class ConfigManager(object):
+    AUTHENTICATED_DATAFILE_URL_TEMPLATE = 'https://config.optimizely.com/datafiles/auth/{sdk_key}.json'
+    AUTHORIZATION_HEADER_DATA_TEMPLATE = 'Bearer {access_token}'
     DATAFILE_URL_TEMPLATE = 'https://cdn.optimizely.com/datafiles/{sdk_key}.json'
     # Default time in seconds to block the 'get_config' method call until 'config' instance has been initialized.
     DEFAULT_BLOCKING_TIMEOUT = 10
@@ -104,6 +106,7 @@ class Errors(object):
 
 
 class HTTPHeaders(object):
+    AUTHORIZATION = 'Authorization'
     IF_MODIFIED_SINCE = 'If-Modified-Since'
     LAST_MODIFIED = 'Last-Modified'
 

diff --git a/optimizely/optimizely.py b/optimizely/optimizely.py
@@ -17,6 +17,7 @@
 from . import event_builder
 from . import exceptions
 from . import logger as _logging
+from .config_manager import AuthDatafilePollingConfigManager
 from .config_manager import PollingConfigManager
 from .config_manager import StaticConfigManager
 from .error_handler import NoOpErrorHandler as noop_error_handler
@@ -43,6 +44,7 @@ def __init__(
         config_manager=None,
         notification_center=None,
         event_processor=None,
+        access_token=None,
     ):
         """ Optimizely init method for managing Custom projects.
 
@@ -65,6 +67,7 @@ def __init__(
                        By default optimizely.event.event_processor.ForwardingEventProcessor is used
                        which simply forwards events to the event dispatcher.
                        To enable event batching configure and use optimizely.event.event_processor.BatchEventProcessor.
+      access_token: Optional string used to fetch authenticated datafile for a secure project environment.
     """
         self.logger_name = '.'.join([__name__, self.__class__.__name__])
         self.is_valid = True
@@ -87,24 +90,24 @@ def __init__(
             self.logger.exception(str(error))
             return
 
+        config_manager_options = {
+            'datafile': datafile,
+            'logger': self.logger,
+            'error_handler': self.error_handler,
+            'notification_center': self.notification_center,
+            'skip_json_validation': skip_json_validation,
+        }
+
         if not self.config_manager:
             if sdk_key:
-                self.config_manager = PollingConfigManager(
-                    sdk_key=sdk_key,
-                    datafile=datafile,
-                    logger=self.logger,
-                    error_handler=self.error_handler,
-                    notification_center=self.notification_center,
-                    skip_json_validation=skip_json_validation,
-                )
+                config_manager_options['sdk_key'] = sdk_key
+                if access_token:
+                    config_manager_options['access_token'] = access_token
+                    self.config_manager = AuthDatafilePollingConfigManager(**config_manager_options)
+                else:
+                    self.config_manager = PollingConfigManager(**config_manager_options)
             else:
-                self.config_manager = StaticConfigManager(
-                    datafile=datafile,
-                    logger=self.logger,
-                    error_handler=self.error_handler,
-                    notification_center=self.notification_center,
-                    skip_json_validation=skip_json_validation,
-                )
+                self.config_manager = StaticConfigManager(**config_manager_options)
 
         self.event_builder = event_builder.EventBuilder()
         self.decision_service = decision_service.DecisionService(self.logger, user_profile_service)

diff --git a/tests/test_config_manager.py b/tests/test_config_manager.py
@@ -365,9 +365,10 @@ def test_set_last_modified(self, _):
 
     def test_fetch_datafile(self, _):
         """ Test that fetch_datafile sets config and last_modified based on response. """
+        sdk_key = 'some_key'
         with mock.patch('optimizely.config_manager.PollingConfigManager.fetch_datafile'):
-            project_config_manager = config_manager.PollingConfigManager(sdk_key='some_key')
-        expected_datafile_url = 'https://cdn.optimizely.com/datafiles/some_key.json'
+            project_config_manager = config_manager.PollingConfigManager(sdk_key=sdk_key)
+        expected_datafile_url = enums.ConfigManager.DATAFILE_URL_TEMPLATE.format(sdk_key=sdk_key)
         test_headers = {'Last-Modified': 'New Time'}
         test_datafile = json.dumps(self.config_dict_with_features)
         test_response = requests.Response()
@@ -397,3 +398,50 @@ def test_is_running(self, _):
         with mock.patch('optimizely.config_manager.PollingConfigManager.fetch_datafile'):
             project_config_manager = config_manager.PollingConfigManager(sdk_key='some_key')
             self.assertTrue(project_config_manager.is_running)
+
+
+@mock.patch('requests.get')
+class AuthDatafilePollingConfigManagerTest(base.BaseTest):
+    def test_init__access_token_none__fails(self, _):
+        """ Test that initialization fails if access_token is None. """
+        self.assertRaisesRegexp(
+            optimizely_exceptions.InvalidInputException,
+            'access_token cannot be empty or None.',
+            config_manager.AuthDatafilePollingConfigManager,
+            access_token=None
+        )
+
+    def test_set_access_token(self, _):
+        """ Test that access_token is properly set as instance variable. """
+        access_token = 'some_token'
+        sdk_key = 'some_key'
+        with mock.patch('optimizely.config_manager.AuthDatafilePollingConfigManager.fetch_datafile'):
+            project_config_manager = config_manager.AuthDatafilePollingConfigManager(
+                access_token=access_token, sdk_key=sdk_key)
+
+        self.assertEqual(access_token, project_config_manager.access_token)
+
+    def test_fetch_datafile(self, _):
+        """ Test that fetch_datafile sets authorization header in request header and sets config based on response. """
+        access_token = 'some_token'
+        sdk_key = 'some_key'
+        with mock.patch('optimizely.config_manager.AuthDatafilePollingConfigManager.fetch_datafile'):
+            project_config_manager = config_manager.AuthDatafilePollingConfigManager(
+                access_token=access_token, sdk_key=sdk_key)
+        expected_datafile_url = enums.ConfigManager.AUTHENTICATED_DATAFILE_URL_TEMPLATE.format(sdk_key=sdk_key)
+        test_datafile = json.dumps(self.config_dict_with_features)
+        test_response = requests.Response()
+        test_response.status_code = 200
+        test_response._content = test_datafile
+
+        # Call fetch_datafile and assert that request was sent with correct authorization header
+        with mock.patch('requests.get', return_value=test_response) as mock_request:
+            project_config_manager.fetch_datafile()
+
+        mock_request.assert_called_once_with(
+            expected_datafile_url,
+            headers={'Authorization': 'Bearer {access_token}'.format(access_token=access_token)},
+            timeout=enums.ConfigManager.REQUEST_TIMEOUT,
+        )
+
+        self.assertIsInstance(project_config_manager.get_config(), project_config.ProjectConfig)
diff --git a/tests/test_optimizely.py b/tests/test_optimizely.py
@@ -252,6 +252,16 @@ def test_init__sdk_key_and_datafile(self):
 
         self.assertIs(type(opt_obj.config_manager), config_manager.PollingConfigManager)
 
+    def test_init__sdk_key_and_access_token(self):
+        """ Test that if both sdk_key and access_token is provided then AuthDatafilePollingConfigManager is used. """
+
+        with mock.patch('optimizely.config_manager.AuthDatafilePollingConfigManager._set_config'), mock.patch(
+            'threading.Thread.start'
+        ):
+            opt_obj = optimizely.Optimizely(access_token='test_access_token', sdk_key='test_sdk_key')
+
+        self.assertIs(type(opt_obj.config_manager), config_manager.AuthDatafilePollingConfigManager)
+
     def test_invalid_json_raises_schema_validation_off(self):
         """ Test that invalid JSON logs error if schema validation is turned off. """