Address review feedback from cursor

Author: keelerm84

ldclient/impl/datasourcev2/streaming.py

@@ -147,6 +147,7 @@ def sync(self, ss: SelectorStore) -> Generator[Update, None, None]:
         self._running = True
         self._connection_attempt_start_time = time()
 
+        envid = None
         for action in self._sse.all:
             if isinstance(action, Fault):
                 # If the SSE client detects the stream has closed, then it will
@@ -165,7 +166,6 @@ def sync(self, ss: SelectorStore) -> Generator[Update, None, None]:
                     break
                 continue
 
-            envid = None
             if isinstance(action, Start) and action.headers is not None:
                 fallback = action.headers.get(_LD_FD_FALLBACK_HEADER) == 'true'
                 envid = action.headers.get(_LD_ENVID_HEADER)

ldclient/testing/impl/datasourcev2/test_polling_synchronizer.py

@@ -20,7 +20,7 @@
     Selector,
     ServerIntent
 )
-from ldclient.impl.util import UnsuccessfulResponseException, _Fail, _Success
+from ldclient.impl.util import UnsuccessfulResponseException, _Fail, _Success, _LD_ENVID_HEADER, _LD_FD_FALLBACK_HEADER
 from ldclient.interfaces import DataSourceErrorKind, DataSourceState
 from ldclient.testing.mock_components import MockSelectorStore
 
@@ -304,3 +304,169 @@ def test_unrecoverable_error_shuts_down():
         assert False, "Expected StopIteration"
     except StopIteration:
         pass
+
+
+def test_envid_from_success_headers():
+    """Test that environment ID is captured from successful polling response headers"""
+    change_set = ChangeSetBuilder.no_changes()
+    headers = {_LD_ENVID_HEADER: 'test-env-polling-123'}
+    polling_result: PollingResult = _Success(value=(change_set, headers))
+
+    synchronizer = PollingDataSource(
+        poll_interval=0.01, requester=ListBasedRequester(results=iter([polling_result]))
+    )
+
+    valid = next(synchronizer.sync(MockSelectorStore(Selector.no_selector())))
+
+    assert valid.state == DataSourceState.VALID
+    assert valid.error is None
+    assert valid.revert_to_fdv1 is False
+    assert valid.environment_id == 'test-env-polling-123'
+
+
+def test_envid_from_success_with_changeset():
+    """Test that environment ID is captured from polling response with actual changes"""
+    builder = ChangeSetBuilder()
+    builder.start(intent=IntentCode.TRANSFER_FULL)
+    builder.add_put(
+        version=100, kind=ObjectKind.FLAG, key="flag-key", obj={"key": "flag-key"}
+    )
+    change_set = builder.finish(selector=Selector(state="p:SOMETHING:300", version=300))
+    headers = {_LD_ENVID_HEADER: 'test-env-456'}
+    polling_result: PollingResult = _Success(value=(change_set, headers))
+
+    synchronizer = PollingDataSource(
+        poll_interval=0.01, requester=ListBasedRequester(results=iter([polling_result]))
+    )
+    valid = next(synchronizer.sync(MockSelectorStore(Selector.no_selector())))
+
+    assert valid.state == DataSourceState.VALID
+    assert valid.environment_id == 'test-env-456'
+    assert valid.change_set is not None
+    assert len(valid.change_set.changes) == 1
+
+
+def test_envid_from_fallback_headers():
+    """Test that environment ID is captured when fallback header is present on success"""
+    change_set = ChangeSetBuilder.no_changes()
+    headers = {
+        _LD_ENVID_HEADER: 'test-env-fallback',
+        _LD_FD_FALLBACK_HEADER: 'true'
+    }
+    polling_result: PollingResult = _Success(value=(change_set, headers))
+
+    synchronizer = PollingDataSource(
+        poll_interval=0.01, requester=ListBasedRequester(results=iter([polling_result]))
+    )
+
+    valid = next(synchronizer.sync(MockSelectorStore(Selector.no_selector())))
+
+    assert valid.state == DataSourceState.VALID
+    assert valid.revert_to_fdv1 is True
+    assert valid.environment_id == 'test-env-fallback'
+
+
+def test_envid_from_error_headers_recoverable():
+    """Test that environment ID is captured from error response headers for recoverable errors"""
+    builder = ChangeSetBuilder()
+    builder.start(intent=IntentCode.TRANSFER_FULL)
+    builder.add_delete(version=101, kind=ObjectKind.FLAG, key="flag-key")
+    change_set = builder.finish(selector=Selector(state="p:SOMETHING:300", version=300))
+    headers_success = {_LD_ENVID_HEADER: 'test-env-success'}
+    polling_result: PollingResult = _Success(value=(change_set, headers_success))
+
+    headers_error = {_LD_ENVID_HEADER: 'test-env-408'}
+    _failure = _Fail(
+        error="error for test",
+        exception=UnsuccessfulResponseException(status=408),
+        headers=headers_error
+    )
+
+    synchronizer = PollingDataSource(
+        poll_interval=0.01,
+        requester=ListBasedRequester(results=iter([_failure, polling_result])),
+    )
+    sync = synchronizer.sync(MockSelectorStore(Selector.no_selector()))
+    interrupted = next(sync)
+    valid = next(sync)
+
+    assert interrupted.state == DataSourceState.INTERRUPTED
+    assert interrupted.environment_id == 'test-env-408'
+    assert interrupted.error is not None
+    assert interrupted.error.status_code == 408
+
+    assert valid.state == DataSourceState.VALID
+    assert valid.environment_id == 'test-env-success'
+
+
+def test_envid_from_error_headers_unrecoverable():
+    """Test that environment ID is captured from error response headers for unrecoverable errors"""
+    headers_error = {_LD_ENVID_HEADER: 'test-env-401'}
+    _failure = _Fail(
+        error="error for test",
+        exception=UnsuccessfulResponseException(status=401),
+        headers=headers_error
+    )
+
+    synchronizer = PollingDataSource(
+        poll_interval=0.01,
+        requester=ListBasedRequester(results=iter([_failure])),
+    )
+    sync = synchronizer.sync(MockSelectorStore(Selector.no_selector()))
+    off = next(sync)
+
+    assert off.state == DataSourceState.OFF
+    assert off.environment_id == 'test-env-401'
+    assert off.error is not None
+    assert off.error.status_code == 401
+
+
+def test_envid_from_error_with_fallback():
+    """Test that environment ID and fallback are captured from error response"""
+    headers_error = {
+        _LD_ENVID_HEADER: 'test-env-503',
+        _LD_FD_FALLBACK_HEADER: 'true'
+    }
+    _failure = _Fail(
+        error="error for test",
+        exception=UnsuccessfulResponseException(status=503),
+        headers=headers_error
+    )
+
+    synchronizer = PollingDataSource(
+        poll_interval=0.01,
+        requester=ListBasedRequester(results=iter([_failure])),
+    )
+    sync = synchronizer.sync(MockSelectorStore(Selector.no_selector()))
+    off = next(sync)
+
+    assert off.state == DataSourceState.OFF
+    assert off.revert_to_fdv1 is True
+    assert off.environment_id == 'test-env-503'
+
+
+def test_envid_from_generic_error_with_headers():
+    """Test that environment ID is captured from generic errors with headers"""
+    builder = ChangeSetBuilder()
+    builder.start(intent=IntentCode.TRANSFER_FULL)
+    change_set = builder.finish(selector=Selector(state="p:SOMETHING:300", version=300))
+    headers_success = {}
+    polling_result: PollingResult = _Success(value=(change_set, headers_success))
+
+    headers_error = {_LD_ENVID_HEADER: 'test-env-generic'}
+    _failure = _Fail(error="generic error for test", headers=headers_error)
+
+    synchronizer = PollingDataSource(
+        poll_interval=0.01,
+        requester=ListBasedRequester(results=iter([_failure, polling_result])),
+    )
+    sync = synchronizer.sync(MockSelectorStore(Selector.no_selector()))
+    interrupted = next(sync)
+    valid = next(sync)
+
+    assert interrupted.state == DataSourceState.INTERRUPTED
+    assert interrupted.environment_id == 'test-env-generic'
+    assert interrupted.error is not None
+    assert interrupted.error.kind == DataSourceErrorKind.NETWORK_ERROR
+
+    assert valid.state == DataSourceState.VALID