diff --git a/tests/testlib/badserverext.py b/tests/testlib/badserverext.py --- a/tests/testlib/badserverext.py +++ b/tests/testlib/badserverext.py @@ -27,6 +27,12 @@ (The value is a list, multiple values can use used to close a series of requests request) +close-after-recv-patterns + If defined, the `close-after-recv-bytes` values only start counting after the + `read` operation that encountered the defined patterns. + (The value is a list, multiple values can use used to close a series of requests + request) + close-after-send-bytes If defined, close the client socket after sending this many bytes. (The value is a list, multiple values can use used to close a series of requests @@ -65,6 +71,11 @@ ) configitem( b'badserver', + b'close-after-recv-patterns', + default=b'', +) +configitem( + b'badserver', b'close-after-send-bytes', default=b'0', ) @@ -84,15 +95,19 @@ def __init__( self, close_after_recv_bytes, + close_after_recv_patterns, close_after_send_bytes, close_after_send_patterns, ): self._all_close_after_recv_bytes = close_after_recv_bytes + self._all_close_after_recv_patterns = close_after_recv_patterns self._all_close_after_send_bytes = close_after_send_bytes self._all_close_after_send_patterns = close_after_send_patterns self.target_recv_bytes = None self.remaining_recv_bytes = None + self.recv_patterns = None + self.recv_data = b'' self.target_send_bytes = None self.remaining_send_bytes = None self.send_pattern = None @@ -107,6 +122,12 @@ self.target_recv_bytes = None self.remaining_recv_bytes = None + self.recv_data = b'' + if self._all_close_after_recv_patterns: + self.recv_pattern = self._all_close_after_recv_patterns.pop(0) + else: + self.recv_pattern = None + if self._all_close_after_send_bytes: self.target_send_bytes = self._all_close_after_send_bytes.pop(0) self.remaining_send_bytes = self.target_send_bytes @@ -124,6 +145,8 @@ """True, if any processing will be needed""" if self.remaining_recv_bytes is not None: return True + if self.recv_pattern is not None: + return True if self.remaining_send_bytes is not None: return True if self.send_pattern is not None: @@ -184,6 +207,7 @@ When the condition are met the socket is closed """ remaining = self.remaining_recv_bytes + pattern = self.recv_pattern orig = object.__getattribute__(obj, '_orig') bmethod = method.encode('ascii') @@ -192,7 +216,7 @@ requested_size = size actual_size = size - if remaining: + if pattern is None and remaining: if size < 0: actual_size = remaining else: @@ -200,7 +224,7 @@ result = func(actual_size) - if remaining: + if pattern is None and remaining: remaining -= len(result) self.remaining_recv_bytes = remaining @@ -216,6 +240,12 @@ msg %= (bmethod, actual_size, requested_repr, len(result), result) obj._writelog(msg) + if pattern is not None: + self.recv_data += result + if pattern.search(self.recv_data): + # start counting bytes starting with the next read + self.recv_pattern = None + if remaining is not None and remaining <= 0: obj._writelog(b'read limit reached; closing socket') obj._cond_close() @@ -363,6 +393,10 @@ b'badserver', b'close-after-recv-bytes' ) all_recv_bytes = process_bytes_config(all_recv_bytes) + all_recv_pattern = self._ui.config( + b'badserver', b'close-after-recv-patterns' + ) + all_recv_pattern = process_pattern_config(all_recv_pattern) all_send_bytes = self._ui.config( b'badserver', b'close-after-send-bytes' ) @@ -373,6 +407,7 @@ all_send_patterns = process_pattern_config(all_send_patterns) self._cond = ConditionTracker( all_recv_bytes, + all_recv_pattern, all_send_bytes, all_send_patterns, )