def test(output, expected, exit_code=0): ret = (output, exit_code) if output is not None else None env = {'VAR': 'val'} with mock.patch("horovod.run.mpi_run.tiny_shell_exec.execute", return_value=ret) as m: implementation = _get_mpi_implementation(env) self.assertEqual(expected, implementation) m.assert_called_once_with('mpirun --version', env)
def test(output, expected, exit_code=0): ret = (output, exit_code) if output is not None else None with mock.patch("horovod.run.mpi_run.tiny_shell_exec.execute", return_value=ret): implementation = _get_mpi_implementation() self.assertEqual(expected, implementation)