no_warning_call

Author: rohitgr7

tests/helpers/utils.py

@@ -14,6 +14,10 @@
 import functools
 import os
 import traceback
+from contextlib import contextmanager
+from typing import Optional
+
+import pytest
 
 from pytorch_lightning import seed_everything
 from pytorch_lightning.callbacks import ModelCheckpoint
@@ -111,3 +115,18 @@ def inner_f(queue, **kwargs):
         assert result == 1, 'expected 1, but returned %s' % result
 
     return wrapper
+
+
+@contextmanager
+def no_warning_call(warning_type, match: Optional[str] = None):
+    with pytest.warns(None) as record:
+        yield
+
+        try:
+            w = record.pop(warning_type)
+            if not ((match and match in str(w.message)) or w):
+                return
+        except AssertionError:
+            # no warning raised
+            return
+        raise AssertionError(f"`{warning_type}` was raised: {w}")

tests/overrides/test_data_parallel.py

@@ -13,6 +13,7 @@
 )
 from pytorch_lightning.trainer.states import RunningStage
 from tests.helpers import BoringModel
+from tests.helpers.utils import no_warning_call
 
 
 @pytest.mark.parametrize("wrapper_class", [
@@ -65,25 +66,21 @@ def test_lightning_wrapper_module_warn_none_output(wrapper_class):
 
     pl_module.automatic_optimization = False
 
-    with pytest.warns(None, match="Your training_step returned None") as record:
+    with no_warning_call(UserWarning, match="Your training_step returned None"):
         pl_module.running_stage = RunningStage.TRAINING
         wrapped_module()
-        assert not record
 
-    with pytest.warns(None, match="Your test_step returned None") as record:
+    with no_warning_call(UserWarning, match="Your test_step returned None"):
         pl_module.running_stage = RunningStage.TESTING
         wrapped_module()
-        assert not record
 
-    with pytest.warns(None, match="Your validation_step returned None") as record:
+    with no_warning_call(UserWarning, match="Your validation_step returned None"):
         pl_module.running_stage = RunningStage.EVALUATING
         wrapped_module()
-        assert not record
 
-    with pytest.warns(None, match="Your predict returned None") as record:
+    with no_warning_call(UserWarning, match="Your predict returned None"):
         pl_module.running_stage = RunningStage.PREDICTING
         wrapped_module()
-        assert not record
 
     with pytest.warns(None) as record:
         pl_module.running_stage = None