Add test cases for the requeue scontrol calls

Signed-off-by: Max Ehrlich <[email protected]>

Author: Queuecumber

tests/tests_pytorch/trainer/connectors/test_signal_connector.py

@@ -100,6 +100,36 @@ def test_auto_requeue_custom_signal_flag(auto_requeue, requeue_signal):
     connector.teardown()
 
 
+@mock.patch("pytorch_lightning.trainer.connectors.signal_connector.call", mock.MagicMock(return_value=0))
+@mock.patch("pytorch_lightning.trainer.Trainer.save_checkpoint", mock.MagicMock())
+@mock.patch.dict(os.environ, {"SLURM_JOB_ID": "12345"})
+def test_auto_requeue_job():
+    from pytorch_lightning.trainer.connectors.signal_connector import call
+
+    trainer = Trainer(plugins=[SLURMEnvironment()])
+    connector = SignalConnector(trainer)
+    connector.slurm_sigusr_handler_fn(signal.SIGUSR1, None)
+
+    assert call.call_args_list[0].args[0] == ["scontrol", "requeue", "12345"]
+
+    connector.teardown()
+
+
+@mock.patch("pytorch_lightning.trainer.connectors.signal_connector.call", mock.MagicMock(return_value=0))
+@mock.patch("pytorch_lightning.trainer.Trainer.save_checkpoint", mock.MagicMock())
+@mock.patch.dict(os.environ, {"SLURM_JOB_ID": "12345", "SLURM_ARRAY_JOB_ID": "12345", "SLURM_ARRAY_TASK_ID": "1"})
+def test_auto_requeue_array_job():
+    from pytorch_lightning.trainer.connectors.signal_connector import call
+
+    trainer = Trainer(plugins=[SLURMEnvironment()])
+    connector = SignalConnector(trainer)
+    connector.slurm_sigusr_handler_fn(signal.SIGUSR1, None)
+
+    assert call.call_args_list[0].args[0] == ["scontrol", "requeue", "12345_1"]
+
+    connector.teardown()
+
+
 def _registering_signals():
     trainer = Trainer()
     trainer._signal_connector.register_signal_handlers()