Update the test federation client to handle streaming responses (#8130)

Now that the server supports streaming back JSON responses, it would be nice to
show the response as it is streamed, in the test tool.
This commit is contained in:
Richard van der Hoff 2020-08-26 14:11:38 +01:00 committed by GitHub
parent 2e6c90ff84
commit 88b9807ba4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 28 additions and 8 deletions

1
changelog.d/8130.misc Normal file
View File

@ -0,0 +1 @@
Update the test federation client to handle streaming responses.

View File

@ -21,10 +21,12 @@ import argparse
import base64 import base64
import json import json
import sys import sys
from typing import Any, Optional
from urllib import parse as urlparse from urllib import parse as urlparse
import nacl.signing import nacl.signing
import requests import requests
import signedjson.types
import srvlookup import srvlookup
import yaml import yaml
from requests.adapters import HTTPAdapter from requests.adapters import HTTPAdapter
@ -69,7 +71,9 @@ def encode_canonical_json(value):
).encode("UTF-8") ).encode("UTF-8")
def sign_json(json_object, signing_key, signing_name): def sign_json(
json_object: Any, signing_key: signedjson.types.SigningKey, signing_name: str
) -> Any:
signatures = json_object.pop("signatures", {}) signatures = json_object.pop("signatures", {})
unsigned = json_object.pop("unsigned", None) unsigned = json_object.pop("unsigned", None)
@ -122,7 +126,14 @@ def read_signing_keys(stream):
return keys return keys
def request_json(method, origin_name, origin_key, destination, path, content): def request(
method: Optional[str],
origin_name: str,
origin_key: signedjson.types.SigningKey,
destination: str,
path: str,
content: Optional[str],
) -> requests.Response:
if method is None: if method is None:
if content is None: if content is None:
method = "GET" method = "GET"
@ -159,11 +170,14 @@ def request_json(method, origin_name, origin_key, destination, path, content):
if method == "POST": if method == "POST":
headers["Content-Type"] = "application/json" headers["Content-Type"] = "application/json"
result = s.request( return s.request(
method=method, url=dest, headers=headers, verify=False, data=content method=method,
url=dest,
headers=headers,
verify=False,
data=content,
stream=True,
) )
sys.stderr.write("Status Code: %d\n" % (result.status_code,))
return result.json()
def main(): def main():
@ -222,7 +236,7 @@ def main():
with open(args.signing_key_path) as f: with open(args.signing_key_path) as f:
key = read_signing_keys(f)[0] key = read_signing_keys(f)[0]
result = request_json( result = request(
args.method, args.method,
args.server_name, args.server_name,
key, key,
@ -231,7 +245,12 @@ def main():
content=args.body, content=args.body,
) )
json.dump(result, sys.stdout) sys.stderr.write("Status Code: %d\n" % (result.status_code,))
for chunk in result.iter_content():
# we write raw utf8 to stdout.
sys.stdout.buffer.write(chunk)
print("") print("")