test-http-bad-server: factor code dealing with "read" in the new object
authorPierre-Yves David <pierre-yves.david@octobus.net>
Fri, 21 Jan 2022 03:05:43 +0100
changeset 48613 b060e305d79f
parent 48612 11e5cb170d36
child 48614 3efc8644dd00
test-http-bad-server: factor code dealing with "read" in the new object This will make sure both `read` and `readline` do the same processing and make it simpler to update that processing in the future. Differential Revision: https://phab.mercurial-scm.org/D12044
tests/testlib/badserverext.py
--- a/tests/testlib/badserverext.py	Fri Jan 21 00:54:15 2022 +0100
+++ b/tests/testlib/badserverext.py	Fri Jan 21 03:05:43 2022 +0100
@@ -146,6 +146,51 @@
 
         return result
 
+    def forward_read(self, obj, method, size=-1):
+        """call an underlying read function until condition are met
+
+        When the condition are met the socket is closed
+        """
+        remaining = self.remaining_recv_bytes
+
+        orig = object.__getattribute__(obj, '_orig')
+        bmethod = method.encode('ascii')
+        func = getattr(orig, method)
+
+        # No read limit. Call original function.
+        if not remaining:
+            result = func(size)
+            obj._writelog(
+                b'%s(%d) -> (%d) %s' % (bmethod, size, len(result), result)
+            )
+            return result
+
+        origsize = size
+
+        if size < 0:
+            size = remaining
+        else:
+            size = min(remaining, size)
+
+        result = func(size)
+        remaining -= len(result)
+
+        obj._writelog(
+            b'%s(%d from %d) -> (%d) %s'
+            % (bmethod, size, origsize, len(result), result)
+        )
+
+        self.remaining_recv_bytes = remaining
+
+        if remaining <= 0:
+            obj._writelog(b'read limit reached; closing socket')
+            obj._cond_close()
+
+            # This is the easiest way to abort the current request.
+            raise Exception('connection closed after receiving N bytes')
+
+        return result
+
 
 # We can't adjust __class__ on a socket instance. So we define a proxy type.
 class socketproxy(object):
@@ -240,78 +285,12 @@
             self._sock.shutdown(socket.SHUT_RDWR)
 
     def read(self, size=-1):
-        remaining = object.__getattribute__(self, '_cond').remaining_recv_bytes
-
-        # No read limit. Call original function.
-        if not remaining:
-            result = object.__getattribute__(self, '_orig').read(size)
-            self._writelog(
-                b'read(%d) -> (%d) (%s) %s' % (size, len(result), result)
-            )
-            return result
-
-        origsize = size
-
-        if size < 0:
-            size = remaining
-        else:
-            size = min(remaining, size)
-
-        result = object.__getattribute__(self, '_orig').read(size)
-        remaining -= len(result)
-
-        self._writelog(
-            b'read(%d from %d) -> (%d) %s'
-            % (size, origsize, len(result), result)
-        )
-
-        object.__getattribute__(self, '_cond').remaining_recv_bytes = remaining
-
-        if remaining <= 0:
-            self._writelog(b'read limit reached; closing socket')
-            self._close()
-
-            # This is the easiest way to abort the current request.
-            raise Exception('connection closed after receiving N bytes')
-
-        return result
+        cond = object.__getattribute__(self, '_cond')
+        return cond.forward_read(self, 'read', size)
 
     def readline(self, size=-1):
-        remaining = object.__getattribute__(self, '_cond').remaining_recv_bytes
-
-        # No read limit. Call original function.
-        if not remaining:
-            result = object.__getattribute__(self, '_orig').readline(size)
-            self._writelog(
-                b'readline(%d) -> (%d) %s' % (size, len(result), result)
-            )
-            return result
-
-        origsize = size
-
-        if size < 0:
-            size = remaining
-        else:
-            size = min(remaining, size)
-
-        result = object.__getattribute__(self, '_orig').readline(size)
-        remaining -= len(result)
-
-        self._writelog(
-            b'readline(%d from %d) -> (%d) %s'
-            % (size, origsize, len(result), result)
-        )
-
-        object.__getattribute__(self, '_cond').remaining_recv_bytes = remaining
-
-        if remaining <= 0:
-            self._writelog(b'read limit reached; closing socket')
-            self._close()
-
-            # This is the easiest way to abort the current request.
-            raise Exception('connection closed after receiving N bytes')
-
-        return result
+        cond = object.__getattribute__(self, '_cond')
+        return cond.forward_read(self, 'readline', size)
 
     def write(self, data):
         cond = object.__getattribute__(self, '_cond')