fix unit tests

Author: Satish Pasumarthi

test/unit/test_mpi.py

@@ -43,11 +43,11 @@ def __init__(self, *args, **kw):
 
 
 @patch("sagemaker_training.mpi._write_env_vars_to_file")
+@patch("sagemaker_training.mpi.logger")
 @patch("os.path.exists")
 @patch("time.sleep")
 @patch("paramiko.SSHClient", new_callable=MockSSHClient)
 @patch("sagemaker_training.mpi._on_terminate")
-@patch("sagemaker_training.mpi.WorkerRunner._wait_for_status_file")
 @patch("psutil.wait_procs")
 @patch("psutil.process_iter")
 @patch("paramiko.AutoAddPolicy")
@@ -57,18 +57,18 @@ def test_mpi_worker_run(
     policy,
     process_iter,
     wait_procs,
-    wait_for_status_file,
     on_terminate,
     ssh_client,
     sleep,
     path_exists,
+    logger,
     write_env_vars,
 ):
 
     process = MagicMock(info={"name": "orted"})
     process_iter.side_effect = lambda attrs: [process]
     wait_procs.return_value = (process, None)
-    wait_for_status_file.return_value = True
+    path_exists.side_effect = [True, False, True]
     worker = mpi.WorkerRunner(
         user_entry_point="train.sh",
         args=["-v", "--lr", "35"],
@@ -79,7 +79,6 @@ def test_mpi_worker_run(
     )
 
     worker.run()
-
     write_env_vars.assert_called_once()
 
     ssh_client().load_system_host_keys.assert_called()
@@ -90,8 +89,8 @@ def test_mpi_worker_run(
 
     popen.assert_called_with(["/usr/sbin/sshd", "-D"])
     path_exists.call_count == 2
-    wait_for_status_file.assert_called_once()
-    wait_for_status_file.assert_called_with("/tmp/done.algo-1")
+    path_exists.return_value = True
+    logger.info.assert_called_with("MPI process finished.")
 
 
 @patch("sagemaker_training.mpi._write_env_vars_to_file")