changeset 48616:98c502a2c462

test-http-bad-server: refactor the writing logic to avoid early return

Our ultimate goal is to add another way to define the connection needs to be
closed. To do so, we need the "read" code to be more unified.

Differential Revision: https://phab.mercurial-scm.org/D12047

--- a/tests/testlib/badserverext.py	Fri Jan 21 11:58:55 2022 +0100
+++ b/tests/testlib/badserverext.py	Fri Jan 21 12:44:39 2022 +0100
@@ -108,22 +108,25 @@
         orig = object.__getattribute__(obj, '_orig')
         bmethod = method.encode('ascii')
         func = getattr(orig, method)
-        # No byte limit on this operation. Call original function.
+
+        if remaining:
+            remaining = max(0, remaining)
+
         if not remaining:
-            result = func(data, *args, **kwargs)
-            obj._writelog(b'%s(%d) -> %s' % (bmethod, len(data), data))
-            return result
-
-        remaining = max(0, remaining)
-
-        if remaining > 0:
+            newdata = data
+        else:
             if remaining < len(data):
                 newdata = data[0:remaining]
             else:
                 newdata = data
+            remaining -= len(newdata)
+            self.remaining_send_bytes = remaining
 
-            remaining -= len(newdata)
+        result = func(newdata, *args, **kwargs)
 
+        if remaining is None:
+            obj._writelog(b'%s(%d) -> %s' % (bmethod, len(data), data))
+        else:
             obj._writelog(
                 b'%s(%d from %d) -> (%d) %s'
                 % (
@@ -135,11 +138,7 @@
                 )
             )
 
-            result = func(newdata, *args, **kwargs)
-
-        self.remaining_send_bytes = remaining
-
-        if remaining <= 0:
+        if remaining is not None and remaining <= 0:
             obj._writelog(b'write limit reached; closing socket')
             object.__getattribute__(obj, '_cond_close')()
             raise Exception('connection closed after sending N bytes')