Fix race in tests by completing executor tasks before assertion

Author: Kwame Efah

mixpanel/flags/test_local_feature_flags.py

@@ -227,19 +227,22 @@ async def test_get_variant_value_tracks_exposure_when_variant_selected(self):
         with patch('mixpanel.flags.utils.normalized_hash') as mock_hash:
             mock_hash.return_value = 0.5
             _ = flags.get_variant_value("test_flag", "fallback", {"distinct_id": "user123"})
+            flags._executor.shutdown()
             flags._tracker.assert_called_once()
 
     @respx.mock
     async def test_get_variant_value_does_not_track_exposure_on_fallback(self):
         flags = await self.setup_flags([])
         _ = flags.get_variant_value("nonexistent_flag", "fallback", {"distinct_id": "user123"})
+        flags._executor.shutdown()
         flags._tracker.assert_not_called()
 
     @respx.mock
     async def test_get_variant_value_does_not_track_exposure_without_distinct_id(self):
         flag = create_test_flag(context="company")
         flags = await self.setup_flags([flag])
         _ = flags.get_variant_value("nonexistent_flag", "fallback", {"company_id": "company123"})
+        flags._executor.shutdown()
         flags._tracker.assert_not_called()
 
     @respx.mock
@@ -283,7 +286,6 @@ async def track_fetch_calls(self):
 
             flags_in_order=[[flag_v1], [flag_v2]]
             flags = await self.setup_flags_with_polling(flags_in_order)
-
             async with polling_limit_check:
                 await polling_limit_check.wait_for(lambda: polling_iterations >= len(flags_in_order))
 

mixpanel/flags/test_remote_feature_flags.py

@@ -64,12 +64,14 @@ async def test_get_variant_value_tracks_exposure_event_if_variant_selected(self)
         if pending:
             await asyncio.gather(*pending, return_exceptions=True)
 
+        self._flags._executor.shutdown()
         self.mock_tracker.assert_called_once()
 
     @respx.mock
     async def test_get_variant_value_does_not_track_exposure_event_if_fallback(self): 
         respx.get(ENDPOINT).mock(side_effect=httpx.RequestError("Network error"))
         await self._flags.aget_variant_value("test_flag", "control", {"distinct_id": "user123"})
+        self._flags._executor.shutdown()
         self.mock_tracker.assert_not_called()
 
     @respx.mock
@@ -133,13 +135,14 @@ def test_get_variant_value_tracks_exposure_event_if_variant_selected(self):
             return_value=create_success_response({"test_flag": SelectedVariant(variant_key="treatment", variant_value="treatment")}))
 
         self._flags.get_variant_value("test_flag", "control", {"distinct_id": "user123"})
+        self._flags._executor.shutdown()
         self.mock_tracker.assert_called_once()
 
     @respx.mock
     def test_get_variant_value_does_not_track_exposure_event_if_fallback(self):
         respx.get(ENDPOINT).mock(side_effect=httpx.RequestError("Network error"))
-        
         self._flags.get_variant_value("test_flag", "control", {"distinct_id": "user123"})
+        self._flags._executor.shutdown()
         self.mock_tracker.assert_not_called()
 
     @respx.mock