Fix flaky test by flushing the progress bar output upon uninstall.

PiperOrigin-RevId: 674041785

langfun/core/concurrent.py

@@ -469,7 +469,9 @@ def refresh(cls) -> None:
       # Process uninstall requests.
       if cls._uninstall_requests:
         for bar_id in cls._uninstall_requests:
-          cls._progress_bars.pop(bar_id, None)
+          bar = cls._progress_bars.pop(bar_id, None)
+          if bar is not None:
+            bar.close()
         cls._uninstall_requests.clear()
 
 
@@ -765,6 +767,10 @@ def update(self, delta):
   def refresh(self) -> None:
     """Refresh progress bar."""
 
+  @abc.abstractmethod
+  def close(self) -> None:
+    """Close progress bar."""
+
 
 class _TqdmProgressControl(_ProgressControl):
   """Tqdm-based progress control."""
@@ -791,6 +797,9 @@ def refresh(self):
     self._tqdm.colour = self.color
     self._tqdm.refresh()
 
+  def close(self):
+    self._tqdm.close()
+
 
 class _ConsoleProgressControl(_ProgressControl):
   """Simple progress control by printing the status to the console."""
@@ -824,6 +833,9 @@ def refresh(self):
       s.write(f' : {status}')
     sys.stderr.write(s.getvalue() + '\n')
 
+  def close(self):
+    sys.stderr.flush()
+
 
 class _NoopProgressControl(_ProgressControl):
   """No-op progress control."""
@@ -834,6 +846,9 @@ def update(self, delta: int) -> None:
   def refresh(self) -> None:
     pass
 
+  def close(self) -> None:
+    pass
+
 
 def _progress_control(
     total: int,

langfun/core/concurrent_test.py

@@ -576,8 +576,8 @@ def fun(x):
             (3, pg.MISSING_VALUE),
         ],
     )
-    self.assertIn('100%', string_io.getvalue())
     concurrent.ProgressBar.uninstall(bar_id)
+    self.assertIn('100%', string_io.getvalue())
 
 
 class ExecutorPoolTest(unittest.TestCase):