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

[indygreg (Gregory Szorc)]

indygreg created this revision.
Herald added a subscriber: mercurial-devel.
Herald added a reviewer: hg-reviewers.

REVISION SUMMARY
  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.

REPOSITORY
  rHG Mercurial

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
Cc: mercurial-devel