From a31f79ce0bcae00ae85f415c3adb3326755dd284 Mon Sep 17 00:00:00 2001
From: Lu Teng <teng.lu@intel.com>
Date: Tue, 24 Sep 2024 18:34:11 +0800
Subject: [PATCH] Flush stdout buffer before checking.

---
 tests/api_test.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/api_test.py b/tests/api_test.py
index adce61d650d6..a6b3cf61366b 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -10883,6 +10883,8 @@ def test_call_wrapped_second_phase_cleanup(self):
 class EnvironmentInfoTest(jtu.JaxTestCase):
   @parameterized.parameters([True, False])
   def test_print_environment_info(self, return_string):
+    # Flush stdout buffer before checking.
+    sys.stdout.flush()
     with jtu.capture_stdout() as stdout:
       result = jax.print_environment_info(return_string=return_string)
     if return_string: