aws-crt 0.1.5 → 0.1.6

data/aws-crt-ffi/crt/s2n/tests/integrationv2/processes.py

@@ -33,6 +33,7 @@ class _processCommunicator(object):
     framework. We rely on sleeps and waits, and still hit hard to debug deadlocks from
     time to time.
     """
+
     def __init__(self, proc):
         self.proc = proc
         self.wait_for_marker = None
@@ -59,7 +60,8 @@ class _processCommunicator(object):
 
         return (stdout, stderr)
 
-    def communicate(self, input_data=None, send_marker_list=None, close_marker=None, timeout=None):
+    def communicate(self, input_data=None, send_marker_list=None, close_marker=None, kill_marker=None,
+                    send_with_newline=False, timeout=None):
         """
         Communicates with the managed process. If send_marker_list is set, input_data will not be sent
         until the marker is seen.
@@ -72,13 +74,21 @@ class _processCommunicator(object):
         stderr = None
 
         try:
-            stdout, stderr = self._communicate(input_data, send_marker_list, close_marker, timeout)
+            stdout, stderr = self._communicate(
+                input_data,
+                send_marker_list,
+                close_marker,
+                kill_marker,
+                send_with_newline,
+                timeout
+            )
         finally:
             self._communication_started = True
 
         return (stdout, stderr)
 
-    def _communicate(self, input_data=None, send_marker_list=None, close_marker=None, timeout=None):
+    def _communicate(self, input_data=None, send_marker_list=None, close_marker=None, kill_marker=None,
+                     send_with_newline=False, timeout=None):
         """
         This method will read and write data to a subprocess in a non-blocking manner.
         The code is heavily based on Popen.communicate. There are a couple differences:
@@ -147,7 +157,7 @@ class _processCommunicator(object):
                 for key, events in ready:
                     # STDIN is only registered to receive events after the send_marker is found.
                     if key.fileobj is self.proc.stdin:
-                        chunk = input_view[input_data_offset :
+                        chunk = input_view[input_data_offset:
                                            input_data_offset + _PIPE_BUF]
                         try:
                             input_data_offset += os.write(key.fd, chunk)
@@ -174,8 +184,11 @@ class _processCommunicator(object):
                         # just mark input_send as true so we can close out STDIN.
                         if send_marker is not None and send_marker in str(data):
                             if self.proc.stdin and input_data:
-                                selector.register(self.proc.stdin, selectors.EVENT_WRITE)
+                                selector.register(
+                                    self.proc.stdin, selectors.EVENT_WRITE)
                                 message = input_data.pop(0)
+                                if send_with_newline:
+                                    message += b'\n'
                                 # Data destined for stdin is stored in a memoryview
                                 input_view = memoryview(message)
                                 input_data_len = len(message)
@@ -187,6 +200,11 @@ class _processCommunicator(object):
                             selector.unregister(self.proc.stderr)
                             return None, None
 
+                        if kill_marker is not None and kill_marker in data:
+                            selector.unregister(self.proc.stdout)
+                            selector.unregister(self.proc.stderr)
+                            self.proc.kill()
+
                 # If we have finished sending all our input, and have received the
                 # ready-to-send marker, we can close out stdin.
                 if self.proc.stdin and input_data_sent:
@@ -223,9 +241,9 @@ class _processCommunicator(object):
             return
         if skip_check_and_raise or _time() > endtime:
             raise subprocess.TimeoutExpired(
-                    self.proc.args, orig_timeout,
-                    output=b''.join(stdout_seq) if stdout_seq else None,
-                    stderr=b''.join(stderr_seq) if stderr_seq else None)
+                self.proc.args, orig_timeout,
+                output=b''.join(stdout_seq) if stdout_seq else None,
+                stderr=b''.join(stderr_seq) if stderr_seq else None)
 
 
 class ManagedProcess(threading.Thread):
@@ -236,8 +254,10 @@ class ManagedProcess(threading.Thread):
     The stdin/stdout/stderr and exist code a monitored and results
     are made available to the caller.
     """
-    def __init__(self, cmd_line, provider_set_ready_condition, wait_for_marker=None, send_marker_list=None, close_marker=None,
-                 timeout=5, data_source=None, env_overrides=dict(), expect_stderr=False):
+
+    def __init__(self, cmd_line, provider_set_ready_condition, wait_for_marker=None, send_marker_list=None,
+                 close_marker=None, timeout=5, data_source=None, env_overrides=dict(), expect_stderr=False,
+                 kill_marker=None, send_with_newline=False):
         threading.Thread.__init__(self)
 
         proc_env = os.environ.copy()
@@ -269,6 +289,8 @@ class ManagedProcess(threading.Thread):
         self.data_source = data_source
         self.send_marker_list = send_marker_list
         self.expect_stderr = expect_stderr
+        self.kill_marker = kill_marker
+        self.send_with_newline = send_with_newline
 
         if data_source is not None:
             if type(data_source) is not list:
@@ -281,10 +303,12 @@ class ManagedProcess(threading.Thread):
     def run(self):
         with self.results_condition:
             try:
-                proc = subprocess.Popen(self.cmd_line, env=self.proc_env, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, close_fds=True)
+                proc = subprocess.Popen(self.cmd_line, env=self.proc_env, stdin=subprocess.PIPE,
+                                        stdout=subprocess.PIPE, stderr=subprocess.PIPE, close_fds=True)
                 self.proc = proc
             except Exception as ex:
-                self.results = Results(None, None, None, ex, self.expect_stderr)
+                self.results = Results(
+                    None, None, None, ex, self.expect_stderr)
                 raise ex
 
             communicator = _processCommunicator(proc)
@@ -298,25 +322,43 @@ class ManagedProcess(threading.Thread):
 
             proc_results = None
             try:
-                proc_results = communicator.communicate(input_data=self.data_source, send_marker_list=self.send_marker_list, close_marker=self.close_marker, timeout=self.timeout)
-                self.results = Results(proc_results[0], proc_results[1], proc.returncode, None, self.expect_stderr)
+                proc_results = communicator.communicate(
+                    input_data=self.data_source,
+                    send_marker_list=self.send_marker_list,
+                    close_marker=self.close_marker,
+                    kill_marker=self.kill_marker,
+                    send_with_newline=self.send_with_newline,
+                    timeout=self.timeout,
+                )
+                self.results = Results(
+                    proc_results[0],
+                    proc_results[1],
+                    proc.returncode,
+                    None,
+                    expect_stderr=self.expect_stderr,
+                    expect_nonzero_exit=self.kill_marker is not None
+                )
             except subprocess.TimeoutExpired as ex:
                 proc.kill()
                 wrapped_ex = TimeoutException(ex)
 
                 # Read any remaining output
                 proc_results = communicator.communicate()
-                self.results = Results(proc_results[0], proc_results[1], proc.returncode, wrapped_ex, self.expect_stderr)
+                self.results = Results(
+                    proc_results[0], proc_results[1], proc.returncode, wrapped_ex, self.expect_stderr)
             except Exception as ex:
-                self.results = Results(proc_results[0], proc_results[1], proc.returncode, ex, self.expect_stderr)
+                self.results = Results(
+                    proc_results[0], proc_results[1], proc.returncode, ex, self.expect_stderr)
                 raise ex
             finally:
                 # This data is dumped to stdout so we capture this
                 # information no matter where a test fails.
                 print("Command line: {}".format(" ".join(self.cmd_line)))
                 print("Exit code: {}".format(proc.returncode))
-                print("Stdout: {}".format(proc_results[0].decode("utf-8", "backslashreplace")))
-                print("Stderr: {}".format(proc_results[1].decode("utf-8", "backslashreplace")))
+                print("Stdout: {}".format(
+                    proc_results[0].decode("utf-8", "backslashreplace")))
+                print("Stderr: {}".format(
+                    proc_results[1].decode("utf-8", "backslashreplace")))
 
     def kill(self):
         self.proc.kill()
@@ -345,7 +387,8 @@ class ManagedProcess(threading.Thread):
         Return the results, or raise the timeout exception.
         """
         with self.results_condition:
-            result = self.results_condition.wait_for(self._results_ready, timeout=self.timeout)
+            result = self.results_condition.wait_for(
+                self._results_ready, timeout=self.timeout)
 
             if result is False:
                 raise Exception("Timeout")