Refactor get_filename so it can be tested

Just pass in the headers instead of the whole response object so that it
can be tested without needing an actual server.
This commit is contained in:
Brian C. Lane 2018-03-14 14:06:24 -07:00
parent d0676dbf7c
commit d4b99b5ee9
1 changed files with 11 additions and 5 deletions

View File

@ -26,6 +26,12 @@ from composer.unix_socket import UnixHTTPConnectionPool
def api_url(api_version, url):
"""Return the versioned path to the API route
:param api_version: The version of the API to talk to. eg. "0"
:type api_version: str
:param url: The API route to talk to
:type url: str
:returns: The full url to use for the route and API version
:rtype: str
"""
return os.path.normpath("/api/v%s/%s" % (api_version, url))
@ -124,7 +130,7 @@ def post_url_json(socket_path, url, body):
headers={"Content-Type": "application/json"})
return json.loads(r.data.decode("utf-8"))
def get_filename(response):
def get_filename(headers):
"""Get the filename from the response header
:param response: The urllib3 response object
@ -133,12 +139,12 @@ def get_filename(response):
:returns: Filename from content-disposition header
:rtype: str
"""
log.debug("Headers = %s", response.headers)
if "content-disposition" not in response.headers:
log.debug("Headers = %s", headers)
if "content-disposition" not in headers:
raise RuntimeError("No Content-Disposition header; cannot get filename")
try:
k, _, v = response.headers["content-disposition"].split(";")[1].strip().partition("=")
k, _, v = headers["content-disposition"].split(";")[1].strip().partition("=")
if k != "filename":
raise RuntimeError("No filename= found in content-disposition header")
except RuntimeError:
@ -163,7 +169,7 @@ def download_file(socket_path, url, progress=True):
if not err["status"]:
raise RuntimeError(err["error"]["msg"])
filename = get_filename(r)
filename = get_filename(r.headers)
if os.path.exists(filename):
msg = "%s exists, skipping download" % filename
log.error(msg)