use catch_warnings context manager

apache · May 16, 2022 · 0a93326 · 0a93326
1 parent 9920440
commit 0a93326
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/sdks/python/apache_beam/dataframe/frames_test.py b/sdks/python/apache_beam/dataframe/frames_test.py
@@ -1604,8 +1604,9 @@ class GroupByTest(_AbstractFrameTest):
   """Tests for DataFrame/Series GroupBy operations."""
   @staticmethod
   def median_sum_fn(x):
-    warnings.filterwarnings("ignore", message="Mean of empty slice")
-    return (x.foo + x.bar).median()
+    with warnings.catch_warnings():
+      warnings.filterwarnings("ignore", message="Mean of empty slice")
+      return (x.foo + x.bar).median()
 
   @parameterized.expand(ALL_GROUPING_AGGREGATIONS)
   def test_groupby_agg(self, agg_type):

diff --git a/sdks/python/apache_beam/dataframe/transforms_test.py b/sdks/python/apache_beam/dataframe/transforms_test.py
@@ -133,8 +133,9 @@ def test_groupby_apply(self):
     })
 
     def median_sum_fn(x):
-      warnings.filterwarnings("ignore", message="Mean of empty slice")
-      return (x.foo + x.bar).median()
+      with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", message="Mean of empty slice")
+        return (x.foo + x.bar).median()
 
     describe = lambda df: df.describe()