changeset 35781:c6ef8e841873

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

We add the ability to specify arbitrary HTTP request headers and
to save the HTTP response body to a file. These will be used in
upcoming commits.

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

--- a/tests/get-with-headers.py	Sat Jan 20 14:59:08 2018 -0800
+++ b/tests/get-with-headers.py	Sat Jan 20 16:08:07 2018 -0800
@@ -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,6 +73,11 @@
         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:
@@ -71,9 +86,13 @@
             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')