D1921: tests: teach get-with-headers.py some new tricks

[indygreg (Gregory Szorc)]

This revision was automatically updated to reflect the committed changes.
Closed by commit rHGc6ef8e841873: tests: teach get-with-headers.py some new tricks (authored by indygreg, committed by ).

REPOSITORY
  rHG Mercurial

CHANGES SINCE LAST UPDATE
  https://phab.mercurial-scm.org/D1921?vs=4962&id=4979

REVISION DETAIL
  https://phab.mercurial-scm.org/D1921

AFFECTED FILES
  tests/get-with-headers.py

CHANGE DETAILS

diff --git a/tests/get-with-headers.py b/tests/get-with-headers.py
--- a/tests/get-with-headers.py
+++ b/tests/get-with-headers.py
@@ -28,6 +28,11 @@
 parser.add_argument('--headeronly', action='store_true')
 parser.add_argument('--json', action='store_true')
 parser.add_argument('--hgproto')
+parser.add_argument('--requestheader', nargs='*', default=[],
+                    help='Send an additional HTTP request header. Argument '
+                         'value is <header>=<value>')
+parser.add_argument('--bodyfile',
+                    help='Write HTTP response body to a file')
 parser.add_argument('host')
 parser.add_argument('path')
 parser.add_argument('show', nargs='*')
@@ -38,6 +43,7 @@
 headeronly = args.headeronly
 formatjson = args.json
 hgproto = args.hgproto
+requestheaders = args.requestheader
 
 tag = None
 def request(host, path, show):
@@ -49,6 +55,10 @@
     if hgproto:
         headers['X-HgProto-1'] = hgproto
 
+    for header in requestheaders:
+        key, value = header.split('=', 1)
+        headers[key] = value
+
     conn = httplib.HTTPConnection(host)
     conn.request("GET", '/' + path, None, headers)
     response = conn.getresponse()
@@ -63,17 +73,26 @@
         print()
         data = response.read()
 
+        if args.bodyfile:
+            bodyfh = open(args.bodyfile, 'wb')
+        else:
+            bodyfh = sys.stdout
+
         # Pretty print JSON. This also has the beneficial side-effect
         # of verifying emitted JSON is well-formed.
         if formatjson:
             # json.dumps() will print trailing newlines. Eliminate them
             # to make tests easier to write.
             data = json.loads(data)
             lines = json.dumps(data, sort_keys=True, indent=2).splitlines()
             for line in lines:
-                print(line.rstrip())
+                bodyfh.write(line.rstrip())
+                bodyfh.write(b'\n')
         else:
-            sys.stdout.write(data)
+            bodyfh.write(data)
+
+        if args.bodyfile:
+            bodyfh.close()
 
     if twice and response.getheader('ETag', None):
         tag = response.getheader('ETag')



To: indygreg, #hg-reviewers, lothiraldan, durin42
Cc: mercurial-devel