generalize test

Author: awaelchli

tests/overrides/test_data_parallel.py

@@ -3,37 +3,45 @@
 import pytest
 import torch
 
-from pytorch_lightning.overrides.data_parallel import LightningDistributedModule
+from pytorch_lightning.overrides.data_parallel import LightningDistributedModule, LightningParallelModule
 
 
-def test_lightning_distributed_module_methods():
+@pytest.mark.parametrize(['wrapper_class'], [
+    LightningParallelModule,
+    LightningDistributedModule,
+])
+def test_lightning_distributed_module_methods(wrapper_class):
     """ Test that the LightningDistributedModule redirects .forward() to the LightningModule methods. """
     pl_module = MagicMock()
-    dist_module = LightningDistributedModule(pl_module)
+    wrapped_module = wrapper_class(pl_module)
 
     batch = torch.rand(5)
     batch_idx = 3
 
     pl_module.training = True
     pl_module.testing = False
-    dist_module(batch, batch_idx)
+    wrapped_module(batch, batch_idx)
     pl_module.training_step.assert_called_with(batch, batch_idx)
 
     pl_module.training = False
     pl_module.testing = True
-    dist_module(batch, batch_idx)
+    wrapped_module(batch, batch_idx)
     pl_module.test_step.assert_called_with(batch, batch_idx)
 
     pl_module.training = False
     pl_module.testing = False
-    dist_module(batch, batch_idx)
+    wrapped_module(batch, batch_idx)
     pl_module.validation_step.assert_called_with(batch, batch_idx)
 
 
-def test_lightning_distributed_module_warn_none_output():
+@pytest.mark.parametrize(['wrapper_class'], [
+    LightningParallelModule,
+    LightningDistributedModule,
+])
+def test_lightning_distributed_module_warn_none_output(wrapper_class):
     """ Test that the LightningDistributedModule warns about forgotten return statement. """
     pl_module = MagicMock()
-    dist_module = LightningDistributedModule(pl_module)
+    wrapped_module = wrapper_class(pl_module)
 
     pl_module.training_step.return_value = None
     pl_module.validation_step.return_value = None
@@ -42,14 +50,14 @@ def test_lightning_distributed_module_warn_none_output():
     with pytest.warns(UserWarning, match="Your training_step returned None"):
         pl_module.training = True
         pl_module.testing = False
-        dist_module()
+        wrapped_module()
 
     with pytest.warns(UserWarning, match="Your test_step returned None"):
         pl_module.training = False
         pl_module.testing = True
-        dist_module()
+        wrapped_module()
 
     with pytest.warns(UserWarning, match="Your validation_step returned None"):
         pl_module.training = False
         pl_module.testing = False
-        dist_module()
+        wrapped_module()