# HG changeset patch
# User Pierre-Yves David <pierre-yves.david@octobus.net>
# Date 1642765479 -3600
# Node ID 98c502a2c462901d543135139b3d14ccad17ec56
# Parent  e38776a4c2cbe40ad52dfc97c19e2a46e3972d51
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

diff -r e38776a4c2cb -r 98c502a2c462 tests/testlib/badserverext.py
--- 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')