Use subprocess.Popen().wait(timeout=... for timeouts. Related #64 (#65)

Timeout can be float
python-useful-helpers · Aug 21, 2018 · f80f0cb · f80f0cb
1 parent 9dfc2bc
commit f80f0cb
Show file tree

Hide file tree

Showing 7 changed files with 31 additions and 23 deletions.
diff --git a/exec_helpers/__init__.py b/exec_helpers/__init__.py
@@ -46,7 +46,7 @@
     'ExecResult',
 )
 
-__version__ = '2.0.0'
+__version__ = '2.0.1'
 __author__ = "Alexey Stepanov"
 __author_email__ = 'penguinolog@gmail.com'
 __maintainers__ = {

diff --git a/exec_helpers/_ssh_client_base.py b/exec_helpers/_ssh_client_base.py
@@ -654,7 +654,7 @@ def _exec_command(
         :type interface: paramiko.channel.Channel
         :type stdout: typing.Optional[paramiko.ChannelFile]
         :type stderr: typing.Optional[paramiko.ChannelFile]
-        :type timeout: typing.Union[int, None]
+        :type timeout: typing.Union[int, float, None]
         :type verbose: bool
         :param log_mask_re: regex lookup rule to mask command for logger.
                             all MATCHED groups will be replaced by '<*masked*>'
@@ -744,7 +744,7 @@ def execute_through_host(
         auth: typing.Optional[ssh_auth.SSHAuth] = None,
         target_port: int = 22,
         verbose: bool = False,
-        timeout: typing.Union[int, None] = constants.DEFAULT_TIMEOUT,
+        timeout: typing.Union[int, float, None] = constants.DEFAULT_TIMEOUT,
         get_pty: bool = False,
         **kwargs: typing.Dict
     ) -> exec_result.ExecResult:
@@ -761,7 +761,7 @@ def execute_through_host(
         :param verbose: Produce log.info records for command call and output
         :type verbose: bool
         :param timeout: Timeout for command execution.
-        :type timeout: typing.Union[int, None]
+        :type timeout: typing.Union[int, float, None]
         :param get_pty: open PTY on target machine
         :type get_pty: bool
         :rtype: ExecResult
@@ -822,7 +822,7 @@ def execute_together(
         cls,
         remotes: typing.Iterable['SSHClientBase'],
         command: str,
-        timeout: typing.Union[int, None] = constants.DEFAULT_TIMEOUT,
+        timeout: typing.Union[int, float, None] = constants.DEFAULT_TIMEOUT,
         expected: typing.Optional[typing.Iterable[int]] = None,
         raise_on_err: bool = True,
         **kwargs: typing.Dict
@@ -834,7 +834,7 @@ def execute_together(
         :param command: Command for execution
         :type command: str
         :param timeout: Timeout for command execution.
-        :type timeout: typing.Union[int, None]
+        :type timeout: typing.Union[int, float, None]
         :param expected: expected return codes (0 by default)
         :type expected: typing.Optional[typing.Iterable[]]
         :param raise_on_err: Raise exception on unexpected return code

diff --git a/exec_helpers/api.py b/exec_helpers/api.py
@@ -197,7 +197,7 @@ def execute(
         self,
         command: str,
         verbose: bool = False,
-        timeout: typing.Union[int, None] = constants.DEFAULT_TIMEOUT,
+        timeout: typing.Union[int, float, None] = constants.DEFAULT_TIMEOUT,
         **kwargs: typing.Dict
     ) -> exec_result.ExecResult:
         """Execute command and wait for return code.
@@ -247,7 +247,7 @@ def check_call(
         self,
         command: str,
         verbose: bool = False,
-        timeout: typing.Union[int, None] = constants.DEFAULT_TIMEOUT,
+        timeout: typing.Union[int, float, None] = constants.DEFAULT_TIMEOUT,
         error_info: typing.Optional[str] = None,
         expected: typing.Optional[typing.Iterable[typing.Union[int, proc_enums.ExitCodes]]] = None,
         raise_on_err: bool = True,
@@ -297,7 +297,7 @@ def check_stderr(
         self,
         command: str,
         verbose: bool = False,
-        timeout: typing.Union[int, None] = constants.DEFAULT_TIMEOUT,
+        timeout: typing.Union[int, float, None] = constants.DEFAULT_TIMEOUT,
         error_info: typing.Optional[str] = None,
         raise_on_err: bool = True,
         **kwargs: typing.Dict

diff --git a/exec_helpers/exceptions.py b/exec_helpers/exceptions.py
@@ -19,7 +19,7 @@
 from exec_helpers import proc_enums
 from exec_helpers import _log_templates
 
-if typing.TYPE_CHECKING:
+if typing.TYPE_CHECKING:  # pragma: no cover
     from exec_helpers import exec_result  # noqa: F401  # pylint: disable=cyclic-import
 
 __all__ = (

diff --git a/exec_helpers/exec_result.py b/exec_helpers/exec_result.py
@@ -329,7 +329,7 @@ def exit_code(self, new_val: typing.Union[int, proc_enums.ExitCodes]) -> None:
         if self.timestamp:
             raise RuntimeError('Exit code is already received.')
         if not isinstance(new_val, int):
-            raise TypeError('Exit code is strictly int')
+            raise TypeError('Exit code is strictly int, received: {code!r}'.format(code=new_val))
         with self.lock:
             self.__exit_code = proc_enums.exit_code_to_enum(new_val)
             if self.__exit_code != proc_enums.ExitCodes.EX_INVALID:

diff --git a/exec_helpers/subprocess_runner.py b/exec_helpers/subprocess_runner.py
@@ -104,7 +104,7 @@ def _exec_command(
         :param stderr: STDERR pipe or file-like object
         :type stderr: typing.Any
         :param timeout: Timeout for command execution
-        :type timeout: typing.Union[int, None]
+        :type timeout: typing.Union[int, float, None]
         :param verbose: produce verbose log record on command call
         :type verbose: bool
         :param log_mask_re: regex lookup rule to mask command for logger.
@@ -143,8 +143,12 @@ def poll_stderr() -> None:
         stderr_future = poll_stderr()  # type: concurrent.futures.Future
         # pylint: enable=assignment-from-no-return
 
-        concurrent.futures.wait([stdout_future, stderr_future], timeout=timeout)  # Wait real timeout here
-        exit_code = interface.poll()  # Update exit code
+        try:
+            exit_code = interface.wait(timeout=timeout)  # Wait real timeout here, it's python 3 only feature
+        except subprocess.TimeoutExpired:
+            exit_code = interface.poll()  # Update exit code
+
+        concurrent.futures.wait([stdout_future, stderr_future], timeout=1)  # Minimal timeout to complete polling
 
         # Process closed?
         if exit_code is not None:
@@ -183,7 +187,7 @@ def execute_async(  # pylint: disable=signature-differs
         **kwargs: typing.Dict
     ) -> typing.Tuple[subprocess.Popen, None, None, None]:
         """Overload: with stdin."""
-        pass
+        pass  # pragma: no cover
 
     @typing.overload  # noqa: F811
     def execute_async(
@@ -197,7 +201,7 @@ def execute_async(
         **kwargs: typing.Dict
     ) -> typing.Tuple[subprocess.Popen, None, typing.IO, typing.IO]:
         """Overload: no stdin."""
-        pass
+        pass  # pragma: no cover
 
     # pylint: enable=unused-argument
     def execute_async(  # type: ignore  # noqa: F811

diff --git a/test/test_subprocess_runner.py b/test/test_subprocess_runner.py
@@ -36,6 +36,7 @@
 stdout_list = [b' \n', b'2\n', b'3\n', b' \n']
 stderr_list = [b' \n', b'0\n', b'1\n', b' \n']
 print_stdin = 'read line; echo "$line"'
+default_timeout = 60 * 60  # 1 hour
 
 
 class FakeFileStream(object):
@@ -88,6 +89,7 @@ def prepare_close(
         else:
             popen_obj.configure_mock(stderr=None)
         popen_obj.attach_mock(mock.Mock(return_value=ec), 'poll')
+        popen_obj.attach_mock(mock.Mock(return_value=ec), 'wait')
         popen_obj.configure_mock(returncode=ec)
 
         popen.return_value = popen_obj
@@ -140,7 +142,7 @@ def test_001_call(
             mock.call.log(level=logging.DEBUG, msg=self.gen_cmd_result_log_message(result)),
         )
         self.assertIn(
-            mock.call.poll(), popen_obj.mock_calls
+            mock.call.wait(timeout=default_timeout), popen_obj.mock_calls
         )
 
     def test_002_call_verbose(
@@ -198,6 +200,7 @@ def test_004_execute_timeout_fail(
     ):
         popen_obj, exp_result = self.prepare_close(popen)
         popen_obj.poll.return_value = None
+        popen_obj.wait.return_value = None
 
         runner = exec_helpers.Subprocess()
 
@@ -263,7 +266,7 @@ def test_005_execute_no_stdout(
                     msg=self.gen_cmd_result_log_message(result)),
             ])
         self.assertIn(
-            mock.call.poll(), popen_obj.mock_calls
+            mock.call.wait(timeout=default_timeout), popen_obj.mock_calls
         )
 
     def test_006_execute_no_stderr(
@@ -305,7 +308,7 @@ def test_006_execute_no_stderr(
                     msg=self.gen_cmd_result_log_message(result)),
             ])
         self.assertIn(
-            mock.call.poll(), popen_obj.mock_calls
+            mock.call.wait(timeout=default_timeout), popen_obj.mock_calls
         )
 
     def test_007_execute_no_stdout_stderr(
@@ -345,7 +348,7 @@ def test_007_execute_no_stdout_stderr(
                     msg=self.gen_cmd_result_log_message(result)),
             ])
         self.assertIn(
-            mock.call.poll(), popen_obj.mock_calls
+            mock.call.wait(timeout=default_timeout), popen_obj.mock_calls
         )
 
     def test_008_execute_mask_global(
@@ -394,7 +397,7 @@ def test_008_execute_mask_global(
         )
 
         self.assertIn(
-            mock.call.poll(), popen_obj.mock_calls
+            mock.call.wait(timeout=default_timeout), popen_obj.mock_calls
         )
 
     def test_009_execute_mask_local(
@@ -439,7 +442,7 @@ def test_009_execute_mask_local(
             mock.call.log(level=logging.DEBUG, msg=self.gen_cmd_result_log_message(result)),
         )
         self.assertIn(
-            mock.call.poll(), popen_obj.mock_calls
+            mock.call.wait(timeout=default_timeout), popen_obj.mock_calls
         )
 
     def test_004_check_stdin_str(
@@ -767,8 +770,9 @@ def test_013_execute_timeout_done(
 
     ):
         popen_obj, exp_result = self.prepare_close(popen, ec=exec_helpers.ExitCodes.EX_INVALID)
-        popen_obj.poll.side_effect = [None, exec_helpers.ExitCodes.EX_INVALID]
+        popen_obj.poll.return_value = exec_helpers.ExitCodes.EX_INVALID
         popen_obj.attach_mock(mock.Mock(side_effect=OSError), 'kill')
+        popen_obj.wait.return_value = None
 
         runner = exec_helpers.Subprocess()