|
35 | 35 |
|
36 | 36 | from .abstractclient import (AbstractShell, AbstractSFTPClient, |
37 | 37 | AbstractSSHClient, AbstractCommand, |
38 | | - SSHClientException, SFTPFileInfo) |
| 38 | + SSHClientException, SSHClientTimeoutException, |
| 39 | + SFTPFileInfo) |
39 | 40 | from .pythonforward import LocalPortForwarding |
40 | 41 | from .utils import is_bytes, is_list_like, is_unicode, is_truthy |
41 | 42 | from robot.api import logger |
@@ -431,55 +432,78 @@ def _put_file(self, source, destination, mode, newline, path_separator, scp_pres |
431 | 432 | def _get_file(self, remote_path, local_path, scp_preserve_times=False): |
432 | 433 | self._scp_client.get(remote_path, local_path, preserve_times=is_truthy(scp_preserve_times)) |
433 | 434 |
|
| 435 | +class OutputBuffers(): |
| 436 | + def __init__(self, shell): |
| 437 | + self._shell = shell |
| 438 | + self._stdout_filebuffer = shell.makefile('rb', -1) |
| 439 | + self._stderr_filebuffer = shell.makefile_stderr('rb', -1) |
| 440 | + self._stdouts = [] |
| 441 | + self._stderrs = [] |
| 442 | + |
| 443 | + def read(self, do_logging=False): |
| 444 | + if self._shell.recv_ready(): |
| 445 | + stdout_output = self._stdout_filebuffer.read(len(self._shell.in_buffer)) |
| 446 | + if do_logging: |
| 447 | + logger.console(stdout_output) |
| 448 | + self._stdouts.append(stdout_output) |
| 449 | + if self._shell.recv_stderr_ready(): |
| 450 | + stderr_output = self._stderr_filebuffer.read(len(self._shell.in_stderr_buffer)) |
| 451 | + if do_logging: |
| 452 | + logger.console(stderr_output) |
| 453 | + self._stderrs.append(stderr_output) |
| 454 | + |
| 455 | + def serialize(self, encoding): |
| 456 | + stdout = (b''.join(self._stdouts) + self._stdout_filebuffer.read()).decode(encoding) |
| 457 | + stderr = (b''.join(self._stderrs) + self._stderr_filebuffer.read()).decode(encoding) |
| 458 | + return stdout, stderr |
434 | 459 |
|
435 | 460 | class RemoteCommand(AbstractCommand): |
| 461 | + _buffers = None |
436 | 462 |
|
437 | 463 | def read_outputs(self, timeout=None, output_during_execution=False, output_if_timeout=False): |
438 | 464 | stderr, stdout = self._receive_stdout_and_stderr(timeout, output_during_execution, output_if_timeout) |
439 | 465 | rc = self._shell.recv_exit_status() |
440 | 466 | self._shell.close() |
441 | 467 | return stdout, stderr, rc |
442 | 468 |
|
| 469 | + def read_unfinished_outputs(self): |
| 470 | + if self._buffers is None: |
| 471 | + return None, None |
| 472 | + if not self._shell.closed: |
| 473 | + self._buffers.read() # Get remaining unread output from channel |
| 474 | + |
| 475 | + stdout, stderr = self._buffers.serialize(self._encoding) |
| 476 | + self._cleanup_buffers() |
| 477 | + |
| 478 | + return stdout, stderr |
| 479 | + |
| 480 | + def _cleanup_buffers(self): |
| 481 | + self._buffers = None |
| 482 | + |
443 | 483 | def _receive_stdout_and_stderr(self, timeout=None, output_during_execution=False, output_if_timeout=False): |
444 | | - stdout_filebuffer = self._shell.makefile('rb', -1) |
445 | | - stderr_filebuffer = self._shell.makefile_stderr('rb', -1) |
446 | | - stdouts = [] |
447 | | - stderrs = [] |
| 484 | + self._buffers = OutputBuffers(self._shell) |
448 | 485 | while self._shell_open(): |
449 | | - self._flush_stdout_and_stderr(stderr_filebuffer, stderrs, stdout_filebuffer, stdouts, timeout, |
450 | | - output_during_execution, output_if_timeout) |
| 486 | + self._flush_stdout_and_stderr(timeout, output_during_execution, output_if_timeout) |
451 | 487 | time.sleep(0.01) # lets not be so busy |
452 | | - stdout = (b''.join(stdouts) + stdout_filebuffer.read()).decode(self._encoding) |
453 | | - stderr = (b''.join(stderrs) + stderr_filebuffer.read()).decode(self._encoding) |
| 488 | + |
| 489 | + stdout, stderr = self._buffers.serialize(self._encoding) |
| 490 | + |
| 491 | + self._buffers = None |
| 492 | + |
454 | 493 | return stderr, stdout |
455 | 494 |
|
456 | | - def _flush_stdout_and_stderr(self, stderr_filebuffer, stderrs, stdout_filebuffer, stdouts, timeout=None, |
457 | | - output_during_execution=False, output_if_timeout=False): |
| 495 | + def _flush_stdout_and_stderr(self, timeout=None, output_during_execution=False, output_if_timeout=False): |
| 496 | + do_logging = is_truthy(output_during_execution) |
458 | 497 | if timeout: |
459 | 498 | end_time = time.time() + timeout |
460 | 499 | while time.time() < end_time: |
461 | 500 | if self._shell.status_event.wait(0): |
462 | 501 | break |
463 | | - self._output_logging(stderr_filebuffer, stderrs, stdout_filebuffer, stdouts, output_during_execution) |
| 502 | + self._buffers.read(do_logging) |
464 | 503 | if not self._shell.status_event.isSet(): |
465 | | - if is_truthy(output_if_timeout): |
466 | | - logger.info(stdouts) |
467 | | - logger.info(stderrs) |
468 | | - raise SSHClientException('Timed out in %s seconds' % int(timeout)) |
| 504 | + raise SSHClientTimeoutException('Timed out in %s seconds' % int(timeout)) |
469 | 505 | else: |
470 | | - self._output_logging(stderr_filebuffer, stderrs, stdout_filebuffer, stdouts, output_during_execution) |
471 | | - |
472 | | - def _output_logging(self, stderr_filebuffer, stderrs, stdout_filebuffer, stdouts, output_during_execution=False): |
473 | | - if self._shell.recv_ready(): |
474 | | - stdout_output = stdout_filebuffer.read(len(self._shell.in_buffer)) |
475 | | - if is_truthy(output_during_execution): |
476 | | - logger.console(stdout_output) |
477 | | - stdouts.append(stdout_output) |
478 | | - if self._shell.recv_stderr_ready(): |
479 | | - stderr_output = stderr_filebuffer.read(len(self._shell.in_stderr_buffer)) |
480 | | - if is_truthy(output_during_execution): |
481 | | - logger.console(stderr_output) |
482 | | - stderrs.append(stderr_output) |
| 506 | + self._buffers.read(do_logging) |
483 | 507 |
|
484 | 508 | def _shell_open(self): |
485 | 509 | return not (self._shell.closed or |
|
0 commit comments