fix: Only record environment ID from successful responses

Co-Authored-By: rlamb@launchdarkly.com <4955475+kinyoklion@users.noreply.github.com>

Author: devin-ai-integration[bot]

ldclient/impl/datasource/streaming.py

@@ -94,8 +94,6 @@ def run(self):
                         log.info("StreamingUpdateProcessor initialized ok.")
                         self._ready.set()
             elif isinstance(action, Fault):
-                record_environment_id(self._data_source_update_sink, action.headers)
-
                 # If the SSE client detects the stream has closed, then it will emit a fault with no-error. We can
                 # ignore this since we want the connection to continue.
                 if action.error is None:

ldclient/impl/datasystem/fdv2.py

@@ -21,12 +21,7 @@
 from ldclient.impl.listeners import Listeners
 from ldclient.impl.repeating_task import RepeatingTask
 from ldclient.impl.rwlock import ReadWriteLock
-from ldclient.impl.util import (
-    _LD_ENVID_HEADER,
-    _LD_FD_FALLBACK_HEADER,
-    _Fail,
-    log
-)
+from ldclient.impl.util import _LD_FD_FALLBACK_HEADER, _Fail, log
 from ldclient.interfaces import (
     DataSourceErrorInfo,
     DataSourceErrorKind,
@@ -222,8 +217,6 @@ def _run_initializers(self, set_on_ready: Event) -> bool:
 
                 if isinstance(basis_result, _Fail):
                     log.warning("Initializer %s failed: %s", initializer.name, basis_result.error)
-                    if basis_result.headers is not None:
-                        self._record_environment_id(basis_result.headers.get(_LD_ENVID_HEADER))
                     # An error response can still carry the FDv1 fallback directive.
                     if basis_result.headers is not None and \
                             basis_result.headers.get(_LD_FD_FALLBACK_HEADER) == 'true':
@@ -422,7 +415,8 @@ def reader(self: 'FDv2'):
                 if self._stop_event.is_set():
                     return ConditionDirective.FALLBACK
 
-                self._record_environment_id(update.environment_id)
+                if update.state == DataSourceState.VALID:
+                    self._record_environment_id(update.environment_id)
 
                 # Handle the update
                 if update.change_set is not None:
@@ -509,7 +503,7 @@ def _persistent_store_outage_recovery(self, data_store_status: DataStoreStatus):
             log.error("Failed to reinitialize data store", exc_info=err)
 
     def _record_environment_id(self, environment_id: Optional[str]):
-        if environment_id is None:
+        if not isinstance(environment_id, str) or environment_id == '':
             return
 
         with self._lock.write():

ldclient/testing/impl/datasource/test_polling_processor.py

@@ -165,7 +165,7 @@ def test_records_environment_id_from_polling_headers():
     sink = DataSourceUpdateSinkImpl(store, Listeners(), Listeners())
     config._data_source_update_sink = sink
     setup_processor(config)
-    ready.wait()
+    assert ready.wait(2)
 
     assert sink.environment_id == 'env-abc-123'
 
@@ -177,6 +177,6 @@ def test_environment_id_is_none_when_requester_provides_no_headers():
     sink = DataSourceUpdateSinkImpl(store, Listeners(), Listeners())
     config._data_source_update_sink = sink
     setup_processor(config)
-    ready.wait()
+    assert ready.wait(2)
 
     assert sink.environment_id is None

ldclient/testing/impl/datasource/test_streaming.py

@@ -472,7 +472,7 @@ def test_environment_id_is_none_when_not_provided():
                 assert sink.environment_id is None
 
 
-def test_records_environment_id_from_error_response_headers():
+def test_does_not_record_environment_id_from_error_response_headers():
     store = InMemoryFeatureStore()
     ready = Event()
 
@@ -487,7 +487,7 @@ def test_records_environment_id_from_error_response_headers():
             with StreamingUpdateProcessor(config, store, ready, None) as sp:
                 sp.start()
                 ready.wait(start_wait)
-                assert sink.environment_id == 'env-from-error'
+                assert sink.environment_id is None
 
 
 def expect_item(store, kind, item):

ldclient/testing/impl/datasystem/test_fdv2_datasystem.py

@@ -786,7 +786,7 @@ def test_environment_id_from_initializer_basis():
     fdv2.stop()
 
 
-def test_environment_id_from_initializer_error_headers():
+def test_environment_id_is_not_recorded_from_initializer_error_headers():
     init = _StaticInitializer(
         "failing-initializer",
         _Fail(error="boom", exception=None, headers={_LD_ENVID_HEADER: 'env-from-error'}),
@@ -801,7 +801,7 @@ def test_environment_id_from_initializer_error_headers():
     fdv2.start(set_on_ready)
     assert set_on_ready.wait(1), "Data system did not become ready in time"
 
-    assert fdv2.environment_id == 'env-from-error'
+    assert fdv2.environment_id is None
     fdv2.stop()
 
 
@@ -826,13 +826,13 @@ def test_environment_id_from_synchronizer_update():
     fdv2.stop()
 
 
-def test_environment_id_is_retained_when_updates_omit_it():
+def test_environment_id_is_not_recorded_from_non_valid_updates():
     sync_mock: Synchronizer = Mock()
     sync_mock.name = "envid-sync"
     sync_mock.stop = Mock()
     sync_mock.sync.return_value = iter([
         Update(state=DataSourceState.VALID, environment_id="env-from-sync"),
-        Update(state=DataSourceState.INTERRUPTED),
+        Update(state=DataSourceState.INTERRUPTED, environment_id="env-from-interrupted"),
     ])
 
     fdv2 = FDv2(