fix: regression to artifact handling in util/taskcluster.py

Author: ahal

src/taskgraph/util/taskcluster.py

@@ -6,9 +6,10 @@
 import copy
 import datetime
 import functools
+import io
 import logging
 import os
-from typing import Dict, List, Union
+from typing import Any, Dict, List, Union
 
 import requests
 import taskcluster_urls as liburls
@@ -69,13 +70,32 @@ def get_taskcluster_client(service: str):
     return getattr(taskcluster, service[0].upper() + service[1:])(options)
 
 
-def _handle_artifact(path, response):
+def _handle_artifact(
+    path: str, response: Union[requests.Response, Dict[str, Any]]
+) -> Any:
+    if isinstance(response, dict):
+        # When taskcluster client returns non-JSON responses, it wraps them in {"response": <Response>}
+        if "response" in response and isinstance(
+            response["response"], requests.Response
+        ):
+            response = response["response"]
+        else:
+            # If we already a dict (parsed JSON), return it directly.
+            return response
+
+    # We have a response object, load the content based on the path extension.
     if path.endswith(".json"):
         return response.json()
 
     if path.endswith(".yml"):
         return yaml.load_stream(response.content)
 
+    # Otherwise return raw response content
+    if hasattr(response.raw, "tell") and response.raw.tell() > 0:
+        # Stream was already read, create a new one from content. This can
+        # happen when mocking with responses.
+        return io.BytesIO(response.content)
+
     response.raw.read = functools.partial(response.raw.read, decode_content=True)
     return response.raw
 

test/test_util_taskcluster.py

@@ -84,7 +84,6 @@ def test_get_artifact_url(monkeypatch):
     assert tc.get_artifact_url(task_id, path) == expected_proxy
 
 
-@pytest.mark.xfail
 def test_get_artifact(responses, root_url):
     tid = "abc123"
     tc.get_taskcluster_client.cache_clear()
@@ -93,7 +92,6 @@ def test_get_artifact(responses, root_url):
     responses.get(
         f"{root_url}/api/queue/v1/task/{tid}/artifacts/artifact.txt",
         body=b"foobar",
-        stream=True,
     )
     raw = tc.get_artifact(tid, "artifact.txt")
     assert raw.read() == b"foobar"
@@ -194,7 +192,6 @@ def test_find_task_id_batched(responses, root_url):
     assert result == {"index.abc": "abc", "index.def": "def"}
 
 
-@pytest.mark.xfail
 def test_get_artifact_from_index(responses, root_url):
     index = "foo"
     path = "file.txt"
@@ -203,7 +200,6 @@ def test_get_artifact_from_index(responses, root_url):
     responses.get(
         f"{root_url}/api/index/v1/task/{index}/artifacts/{path}",
         body=b"foobar",
-        stream=True,
     )
 
     result = tc.get_artifact_from_index(index, path)