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
=') +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')