no_warning_call

Author: carmocca

tests/deprecated_api/__init__.py

@@ -13,24 +13,9 @@
 # limitations under the License.
 """Test deprecated functionality which will be removed in vX.Y.Z"""
 import sys
-from contextlib import contextmanager
-
-import pytest
 
 
 def _soft_unimport_module(str_module):
     # once the module is imported  e.g with parsing with pytest it lives in memory
     if str_module in sys.modules:
         del sys.modules[str_module]
-
-
-@contextmanager
-def no_deprecated_call():
-    with pytest.warns(None) as record:
-        yield
-        try:
-            w = record.pop(DeprecationWarning)
-        except AssertionError:
-            # no DeprecationWarning raised
-            return
-        raise AssertionError(f"`DeprecationWarning` was raised: {w}")

tests/deprecated_api/test_remove_1-5.py

@@ -16,8 +16,8 @@
 import pytest
 
 from pytorch_lightning import Trainer, Callback
-from tests.deprecated_api import no_deprecated_call
 from tests.helpers import BoringModel
+from tests.helpers.utils import no_warning_call
 
 
 def test_v1_5_0_old_callback_on_save_checkpoint(tmpdir):
@@ -52,5 +52,5 @@ def on_save_checkpoint(self, *args):
             ...
 
     trainer.callbacks = [NewSignature(), ValidSignature1(), ValidSignature2()]
-    with no_deprecated_call():
+    with no_warning_call(DeprecationWarning):
         trainer.save_checkpoint(filepath)

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 w.text) or w):
+                return
+        except AssertionError:
+            # no warning raised
+            return
+        raise AssertionError(f"`{warning_type}` was raised: {w}")