Mercurial > public > mercurial-scm > hg

comparison tests/get-with-headers.py @ 35781:c6ef8e841873

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
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
author Gregory Szorc <gregory.szorc@gmail.com>
date Sat, 20 Jan 2018 16:08:07 -0800
parents 32317f8bbe2a
children c95c8ab2e7ec
comparison
equal deleted inserted replaced
35780:32317f8bbe2a 35781:c6ef8e841873
26 parser = argparse.ArgumentParser() 26 parser = argparse.ArgumentParser()
27 parser.add_argument('--twice', action='store_true') 27 parser.add_argument('--twice', action='store_true')
28 parser.add_argument('--headeronly', action='store_true') 28 parser.add_argument('--headeronly', action='store_true')
29 parser.add_argument('--json', action='store_true') 29 parser.add_argument('--json', action='store_true')
30 parser.add_argument('--hgproto') 30 parser.add_argument('--hgproto')
31 parser.add_argument('--requestheader', nargs='*', default=[],
32 help='Send an additional HTTP request header. Argument '
33 'value is <header>=<value>')
34 parser.add_argument('--bodyfile',
35 help='Write HTTP response body to a file')
31 parser.add_argument('host') 36 parser.add_argument('host')
32 parser.add_argument('path') 37 parser.add_argument('path')
33 parser.add_argument('show', nargs='*') 38 parser.add_argument('show', nargs='*')
34 39
35 args = parser.parse_args() 40 args = parser.parse_args()
36 41
37 twice = args.twice 42 twice = args.twice
38 headeronly = args.headeronly 43 headeronly = args.headeronly
39 formatjson = args.json 44 formatjson = args.json
40 hgproto = args.hgproto 45 hgproto = args.hgproto
46 requestheaders = args.requestheader
41 47
42 tag = None 48 tag = None
43 def request(host, path, show): 49 def request(host, path, show):
44 assert not path.startswith('/'), path 50 assert not path.startswith('/'), path
45 global tag 51 global tag
46 headers = {} 52 headers = {}
47 if tag: 53 if tag:
48 headers['If-None-Match'] = tag 54 headers['If-None-Match'] = tag
49 if hgproto: 55 if hgproto:
50 headers['X-HgProto-1'] = hgproto 56 headers['X-HgProto-1'] = hgproto
57
58 for header in requestheaders:
59 key, value = header.split('=', 1)
60 headers[key] = value
51 61
52 conn = httplib.HTTPConnection(host) 62 conn = httplib.HTTPConnection(host)
53 conn.request("GET", '/' + path, None, headers) 63 conn.request("GET", '/' + path, None, headers)
54 response = conn.getresponse() 64 response = conn.getresponse()
55 print(response.status, response.reason) 65 print(response.status, response.reason)
61 print("%s: %s" % (h, response.getheader(h))) 71 print("%s: %s" % (h, response.getheader(h)))
62 if not headeronly: 72 if not headeronly:
63 print() 73 print()
64 data = response.read() 74 data = response.read()
65 75
76 if args.bodyfile:
77 bodyfh = open(args.bodyfile, 'wb')
78 else:
79 bodyfh = sys.stdout
80
66 # Pretty print JSON. This also has the beneficial side-effect 81 # Pretty print JSON. This also has the beneficial side-effect
67 # of verifying emitted JSON is well-formed. 82 # of verifying emitted JSON is well-formed.
68 if formatjson: 83 if formatjson:
69 # json.dumps() will print trailing newlines. Eliminate them 84 # json.dumps() will print trailing newlines. Eliminate them
70 # to make tests easier to write. 85 # to make tests easier to write.
71 data = json.loads(data) 86 data = json.loads(data)
72 lines = json.dumps(data, sort_keys=True, indent=2).splitlines() 87 lines = json.dumps(data, sort_keys=True, indent=2).splitlines()
73 for line in lines: 88 for line in lines:
74 print(line.rstrip()) 89 bodyfh.write(line.rstrip())
90 bodyfh.write(b'\n')
75 else: 91 else:
76 sys.stdout.write(data) 92 bodyfh.write(data)
93
94 if args.bodyfile:
95 bodyfh.close()
77 96
78 if twice and response.getheader('ETag', None): 97 if twice and response.getheader('ETag', None):
79 tag = response.getheader('ETag') 98 tag = response.getheader('ETag')
80 99
81 return response.status 100 return response.status