Stabilize FFE exposure system tests (#7168)

Co-authored-by: Codex GPT-5 <noreply@openai.com>

Author: leoromanovsky

tests/ffe/test_exposures.py

@@ -14,6 +14,71 @@
 
 RC_PRODUCT = "FFE_FLAGS"
 RC_PATH = f"datadog/2/{RC_PRODUCT}"
+EXPOSURES_PATH = "/api/v2/exposures"
+EXPOSURE_WAIT_TIMEOUT_SECONDS = 30
+
+
+def exposure_events_from_data(
+    data: dict, flag_keys: set[str] | None = None, subject_id: str | None = None
+) -> list[dict]:
+    """Return exposure events from one agent payload matching the optional flag/subject filters."""
+    if data.get("path") != EXPOSURES_PATH:
+        return []
+
+    exposure_data = data.get("request", {}).get("content")
+    if not isinstance(exposure_data, dict):
+        return []
+
+    exposures = exposure_data.get("exposures")
+    if not isinstance(exposures, list):
+        return []
+
+    events = []
+    for event in exposures:
+        if not isinstance(event, dict):
+            continue
+
+        flag = event.get("flag")
+        subject = event.get("subject")
+        event_flag_key = flag.get("key") if isinstance(flag, dict) else None
+        event_subject_id = subject.get("id") if isinstance(subject, dict) else None
+
+        if flag_keys is not None and event_flag_key not in flag_keys:
+            continue
+        if subject_id is not None and event_subject_id != subject_id:
+            continue
+        events.append(event)
+    return events
+
+
+def find_exposure_events(flag_key: str, subject_id: str | None = None) -> list[dict]:
+    """Find captured exposure events for a specific flag key and optionally a specific subject."""
+    events = []
+    for data in interfaces.agent.get_data(path_filters=EXPOSURES_PATH):
+        events.extend(exposure_events_from_data(data, {flag_key}, subject_id))
+    return events
+
+
+def wait_for_exposure_event(flag_keys: set[str], subject_id: str | None = None) -> None:
+    """Wait until the agent receives an exposure event for one of the given flags."""
+    assert interfaces.agent.wait_for(
+        lambda data: bool(exposure_events_from_data(data, flag_keys, subject_id)),
+        timeout=EXPOSURE_WAIT_TIMEOUT_SECONDS,
+    ), f"Timed out waiting for exposure event for flags {sorted(flag_keys)} and subject {subject_id!r}"
+
+
+def wait_for_min_exposure_count(flag_key: str, expected: int, subject_id: str | None = None) -> int:
+    """Wait until enough matching exposure events are available, then return the current count."""
+    count = count_exposure_events(flag_key, subject_id)
+
+    if count < expected:
+        assert interfaces.agent.wait_for(
+            lambda _: count_exposure_events(flag_key, subject_id) >= expected,
+            timeout=EXPOSURE_WAIT_TIMEOUT_SECONDS,
+        ), f"Timed out waiting for exposure count >= {expected} for flag {flag_key} and subject {subject_id!r}"
+        count = count_exposure_events(flag_key, subject_id)
+
+    return count
 
 
 # Simple UFC fixture for testing with doLog: true
@@ -71,12 +136,13 @@ def setup_ffe_exposure_event_generation(self):
     def test_ffe_exposure_event_generation(self):
         """Test that FFE generates exposure events when flags are evaluated via weblog."""
         assert self.r.status_code == 200, f"Flag evaluation failed: {self.r.text}"
+        wait_for_exposure_event({self.flag}, self.targeting_key)
 
         # Search for our specific flag in all exposure events
         matching_event = None
         context_validated = False
 
-        for data in interfaces.agent.get_data(path_filters="/api/v2/exposures"):
+        for data in interfaces.agent.get_data(path_filters=EXPOSURES_PATH):
             # validate data sent to /api/v2/exposures
 
             exposure_data = data["request"]["content"]
@@ -216,11 +282,12 @@ def test_ffe_multiple_remote_config_files(self):
         """Test that FFE correctly handles multiple remote config files with different flags."""
         assert self.r1.status_code == 200, f"First flag evaluation failed: {self.r1.text}"
         assert self.r2.status_code == 200, f"Second flag evaluation failed: {self.r2.text}"
+        wait_for_exposure_event({self.flag_1, self.flag_2}, self.targeting_key)
 
         # Collect all exposure events for our specific flags
         flags_found = set()
 
-        for data in interfaces.agent.get_data(path_filters="/api/v2/exposures"):
+        for data in interfaces.agent.get_data(path_filters=EXPOSURES_PATH):
             exposure_data = data["request"]["content"]
             assert exposure_data is not None, "No exposure events were sent to agent"
 
@@ -286,7 +353,7 @@ def test_ffe_empty_remote_config(self):
 
         # When no remote config is set, FFE should still work but return default value
         # The exposure events should still be generated based on library configuration
-        for data in interfaces.agent.get_data(path_filters="/api/v2/exposures"):
+        for data in interfaces.agent.get_data(path_filters=EXPOSURES_PATH):
             exposure_data = data["request"]["content"]
             if exposure_data is not None:
                 # Validate that context is still present
@@ -381,12 +448,13 @@ def test_ffe_malformed_remote_config_rejection(self):
         """Test that FFE rejects malformed remote config and preserves the old valid configuration."""
         assert self.r1.status_code == 200, f"First flag evaluation failed: {self.r1.text}"
         assert self.r2.status_code == 200, f"Second flag evaluation failed: {self.r2.text}"
+        wait_for_exposure_event({self.flag}, self.targeting_key)
 
         # Verify that exposure events are still generated for both requests
         # and the flag configuration remained valid despite the malformed update
         events_found = []
 
-        for data in interfaces.agent.get_data(path_filters="/api/v2/exposures"):
+        for data in interfaces.agent.get_data(path_filters=EXPOSURES_PATH):
             exposure_data = data["request"]["content"]
             assert exposure_data is not None, "No exposure events were sent to agent"
 
@@ -430,21 +498,7 @@ def count_exposure_events(flag_key: str, subject_id: str | None = None) -> int:
         Number of matching exposure events found
 
     """
-    count = 0
-    for data in interfaces.agent.get_data(path_filters="/api/v2/exposures"):
-        exposure_data = data["request"]["content"]
-        if exposure_data is None:
-            continue
-
-        exposures = exposure_data.get("exposures", [])
-        for event in exposures:
-            event_flag_key = event.get("flag", {}).get("key")
-            event_subject_id = event.get("subject", {}).get("id")
-
-            if event_flag_key == flag_key:
-                if subject_id is None or event_subject_id == subject_id:
-                    count += 1
-    return count
+    return len(find_exposure_events(flag_key, subject_id))
 
 
 @scenarios.feature_flagging_and_experimentation
@@ -488,7 +542,7 @@ def test_ffe_exposure_caching_same_subject(self):
             assert result["value"] == "value-a", f"Request {i + 1}: expected 'value-a', got '{result['value']}'"
 
         # Count exposure events for this specific subject
-        exposure_count = count_exposure_events(self.flag_key, self.targeting_key)
+        exposure_count = wait_for_min_exposure_count(self.flag_key, 1, self.targeting_key)
 
         # The exposure cache should deduplicate events - we expect exactly 1 exposure
         # for the same (subject, allocation, variant) tuple
@@ -538,6 +592,10 @@ def test_ffe_exposure_caching_different_subjects(self):
             result = json.loads(r.text)
             assert result["value"] == "value-a", f"Request {i + 1}: expected 'value-a', got '{result['value']}'"
 
+        # Wait for each subject to be observed before asserting exact totals.
+        for subject in self.subjects:
+            wait_for_min_exposure_count(self.flag_key, 1, subject)
+
         # Count total exposure events for this flag
         total_exposure_count = count_exposure_events(self.flag_key)
 
@@ -642,7 +700,7 @@ def test_ffe_exposure_caching_allocation_cycle(self):
         # - Exposure #1: default-allocation
         # - Exposure #2: different-allocation (allocation changed)
         # - Exposure #3: default-allocation (allocation changed back)
-        exposure_count = count_exposure_events(self.flag_key, self.targeting_key)
+        exposure_count = wait_for_min_exposure_count(self.flag_key, 3, self.targeting_key)
 
         assert exposure_count == 3, (
             f"Expected exactly 3 exposure events for subject '{self.targeting_key}' "
@@ -737,7 +795,7 @@ def test_ffe_exposure_caching_variant_cycle(self):
         # - Exposure #1: variant-a
         # - Exposure #2: variant-b (variant changed)
         # - Exposure #3: variant-a (variant changed back)
-        exposure_count = count_exposure_events(self.flag_key, self.targeting_key)
+        exposure_count = wait_for_min_exposure_count(self.flag_key, 3, self.targeting_key)
 
         assert exposure_count == 3, (
             f"Expected exactly 3 exposure events for subject '{self.targeting_key}' "
@@ -911,11 +969,12 @@ def test_ffe_exp_5_missing_targeting_key(self):
 
         result = json.loads(self.response.text)
         assert result["value"] == "value-a", f"Expected 'value-a', got '{result['value']}'"
+        wait_for_exposure_event({self.flag_key}, "")
 
         # Search for exposure event with empty subject.id
         matching_event = None
         all_events_for_flag = []  # Collect all events for debugging
-        for data in interfaces.agent.get_data(path_filters="/api/v2/exposures"):
+        for data in interfaces.agent.get_data(path_filters=EXPOSURES_PATH):
             exposure_data = data["request"]["content"]
             if exposure_data is None:
                 continue