Skip to content

[AI-FSSDK] (DO NOT REVIEW) [FSSDK-12273] Add support for custom headers in polling datafile managers #489

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
raju-opti wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[AI-FSSDK] (DO NOT REVIEW) [FSSDK-12273] Add support for custom headers in polling datafile managers #489

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions optimizely/config_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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,
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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()

Expand Down
5 changes: 5 additions & 0 deletions optimizely/optimizely.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
159 changes: 159 additions & 0 deletions tests/test_config_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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)

Loading