diff --git a/optimizely/config_manager.py b/optimizely/config_manager.py index 3dce2741..f2b01505 100644 --- a/optimizely/config_manager.py +++ b/optimizely/config_manager.py @@ -203,6 +203,7 @@ def __init__( notification_center: Optional[NotificationCenter] = None, skip_json_validation: Optional[bool] = False, retries: Optional[int] = 3, + custom_headers: Optional[dict[str, str]] = None, ): """ Initialize config manager. One of sdk_key or datafile has to be set to be able to use. @@ -223,9 +224,12 @@ def __init__( skip_json_validation: Optional boolean param which allows skipping JSON schema validation upon object invocation. By default JSON schema validation will be performed. + custom_headers: Optional dictionary of custom headers to include in datafile fetch requests. + User-provided headers take precedence over SDK internal headers. """ self.retries = retries + self.custom_headers = custom_headers or {} self._config_ready_event = threading.Event() super().__init__( datafile=datafile, @@ -394,6 +398,9 @@ def fetch_datafile(self) -> None: if self.last_modified: request_headers[enums.HTTPHeaders.IF_MODIFIED_SINCE] = self.last_modified + # Merge custom headers, with user-provided headers taking precedence + request_headers.update(self.custom_headers) + try: session = requests.Session() @@ -487,6 +494,9 @@ def fetch_datafile(self) -> None: if self.last_modified: request_headers[enums.HTTPHeaders.IF_MODIFIED_SINCE] = self.last_modified + # Merge custom headers, with user-provided headers taking precedence + request_headers.update(self.custom_headers) + try: session = requests.Session() diff --git a/optimizely/optimizely.py b/optimizely/optimizely.py index 4b014e7f..1cee7ca6 100644 --- a/optimizely/optimizely.py +++ b/optimizely/optimizely.py @@ -76,6 +76,7 @@ def __init__( event_processor_options: Optional[dict[str, Any]] = None, settings: Optional[OptimizelySdkSettings] = None, cmab_service: Optional[DefaultCmabService] = None, + custom_headers: Optional[dict[str, str]] = None, ) -> None: """ Optimizely init method for managing Custom projects. @@ -104,6 +105,8 @@ def __init__( default_decide_options: Optional list of decide options used with the decide APIs. event_processor_options: Optional dict of options to be passed to the default batch event processor. settings: Optional instance of OptimizelySdkSettings for sdk configuration. + custom_headers: Optional dictionary of custom headers to include in datafile fetch requests. + User-provided headers take precedence over SDK internal headers. """ self.logger_name = '.'.join([__name__, self.__class__.__name__]) self.is_valid = True @@ -163,6 +166,8 @@ def __init__( if not self.config_manager: if sdk_key: config_manager_options['sdk_key'] = sdk_key + if custom_headers: + config_manager_options['custom_headers'] = custom_headers if datafile_access_token: config_manager_options['datafile_access_token'] = datafile_access_token self.config_manager = AuthDatafilePollingConfigManager(**config_manager_options) diff --git a/tests/test_config_manager.py b/tests/test_config_manager.py index 1930520e..fcd0d9a6 100644 --- a/tests/test_config_manager.py +++ b/tests/test_config_manager.py @@ -531,6 +531,80 @@ def test_is_running(self, _): project_config_manager.stop() + def test_custom_headers(self, _): + """ Test that custom headers are included in datafile fetch requests. """ + sdk_key = 'some_key' + custom_headers = { + 'X-Custom-Header': 'custom_value', + 'X-Another-Header': 'another_value' + } + + 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() + test_response.status_code = 200 + test_response.headers = test_headers + test_response._content = test_datafile + + with mock.patch('requests.Session.get', return_value=test_response) as mock_request: + project_config_manager = config_manager.PollingConfigManager( + sdk_key=sdk_key, + custom_headers=custom_headers + ) + project_config_manager.stop() + + # Assert that custom headers were included in the request + mock_request.assert_called_once_with( + expected_datafile_url, + headers=custom_headers, + timeout=enums.ConfigManager.REQUEST_TIMEOUT + ) + self.assertEqual(test_headers['Last-Modified'], project_config_manager.last_modified) + self.assertIsInstance(project_config_manager.get_config(), project_config.ProjectConfig) + + def test_custom_headers_override_internal_headers(self, _): + """ Test that custom headers override internal SDK headers. """ + sdk_key = 'some_key' + custom_last_modified = 'Custom Last Modified Time' + custom_headers = { + 'If-Modified-Since': custom_last_modified, + 'X-Custom-Header': 'custom_value' + } + + 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() + test_response.status_code = 200 + test_response.headers = test_headers + test_response._content = test_datafile + + # First request to set last_modified + with mock.patch('requests.Session.get', return_value=test_response): + project_config_manager = config_manager.PollingConfigManager( + sdk_key=sdk_key, + custom_headers=custom_headers + ) + project_config_manager.stop() + + # Second request should use custom header value instead of internal last_modified + with mock.patch('requests.Session.get', return_value=test_response) as mock_request: + project_config_manager._initialize_thread() + project_config_manager.start() + project_config_manager.stop() + + # Assert that custom If-Modified-Since header overrides the internal one + expected_headers = { + 'If-Modified-Since': custom_last_modified, # User's value should be used + 'X-Custom-Header': 'custom_value' + } + mock_request.assert_called_once_with( + expected_datafile_url, + headers=expected_headers, + timeout=enums.ConfigManager.REQUEST_TIMEOUT + ) + @mock.patch('requests.Session.get') class AuthDatafilePollingConfigManagerTest(base.BaseTest): @@ -637,3 +711,88 @@ def test_fetch_datafile__request_exception_raised(self, _): ) self.assertEqual(test_headers['Last-Modified'], project_config_manager.last_modified) self.assertIsInstance(project_config_manager.get_config(), project_config.ProjectConfig) + + def test_custom_headers(self, _): + """ Test that custom headers are included in authenticated datafile fetch requests. """ + datafile_access_token = 'some_token' + sdk_key = 'some_key' + custom_headers = { + 'X-Custom-Header': 'custom_value', + 'X-Another-Header': 'another_value' + } + + with mock.patch('optimizely.config_manager.AuthDatafilePollingConfigManager.fetch_datafile'), mock.patch( + 'optimizely.config_manager.AuthDatafilePollingConfigManager._run' + ): + project_config_manager = config_manager.AuthDatafilePollingConfigManager( + datafile_access_token=datafile_access_token, + sdk_key=sdk_key, + custom_headers=custom_headers + ) + + expected_datafile_url = enums.ConfigManager.AUTHENTICATED_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() + test_response.status_code = 200 + test_response.headers = test_headers + test_response._content = test_datafile + + # Call fetch_datafile and assert that request was sent with both authorization and custom headers + with mock.patch('requests.Session.get', return_value=test_response) as mock_request: + project_config_manager.fetch_datafile() + + expected_headers = { + 'Authorization': f'Bearer {datafile_access_token}', + 'X-Custom-Header': 'custom_value', + 'X-Another-Header': 'another_value' + } + mock_request.assert_called_once_with( + expected_datafile_url, + headers=expected_headers, + timeout=enums.ConfigManager.REQUEST_TIMEOUT, + ) + self.assertIsInstance(project_config_manager.get_config(), project_config.ProjectConfig) + + def test_custom_headers_override_authorization(self, _): + """ Test that custom Authorization header overrides internal SDK authorization header. """ + datafile_access_token = 'some_token' + custom_auth = 'Bearer custom_token' + sdk_key = 'some_key' + custom_headers = { + 'Authorization': custom_auth, + 'X-Custom-Header': 'custom_value' + } + + with mock.patch('optimizely.config_manager.AuthDatafilePollingConfigManager.fetch_datafile'), mock.patch( + 'optimizely.config_manager.AuthDatafilePollingConfigManager._run' + ): + project_config_manager = config_manager.AuthDatafilePollingConfigManager( + datafile_access_token=datafile_access_token, + sdk_key=sdk_key, + custom_headers=custom_headers + ) + + expected_datafile_url = enums.ConfigManager.AUTHENTICATED_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() + test_response.status_code = 200 + test_response.headers = test_headers + test_response._content = test_datafile + + # Call fetch_datafile and assert that custom Authorization header is used + with mock.patch('requests.Session.get', return_value=test_response) as mock_request: + project_config_manager.fetch_datafile() + + expected_headers = { + 'Authorization': custom_auth, # User's custom auth should override + 'X-Custom-Header': 'custom_value' + } + mock_request.assert_called_once_with( + expected_datafile_url, + headers=expected_headers, + timeout=enums.ConfigManager.REQUEST_TIMEOUT, + ) + self.assertIsInstance(project_config_manager.get_config(), project_config.ProjectConfig) +