bot: Use taskcluster client library, fixes #12

bot/code_coverage_bot/taskcluster.py

@@ -1,23 +1,15 @@
 # -*- coding: utf-8 -*-
 import os
-import shutil
 from zipfile import BadZipFile
 from zipfile import is_zipfile
 
-import requests
+import taskcluster
 
 from code_coverage_bot.utils import retry
 from code_coverage_tools.taskcluster import TaskclusterConfig
 
 taskcluster_config = TaskclusterConfig()
 
-index_base = "https://index.taskcluster.net/v1/"
-queue_base = "https://queue.taskcluster.net/v1/"
-
-
-class TaskclusterException(Exception):
-    pass
-
 
 def get_task(branch, revision, platform):
     if platform == "linux":
@@ -33,75 +25,63 @@ def get_task(branch, revision, platform):
         platform_name = "android-api-16-ccov-debug"
         product = "mobile"
     else:
-        raise TaskclusterException("Unsupported platform: %s" % platform)
-
-    r = requests.get(
-        index_base
-        + "task/gecko.v2.{}.revision.{}.{}.{}".format(
-            branch, revision, product, platform_name
-        )
-    )
-    task = r.json()
-    if r.status_code == requests.codes.ok:
-        return task["taskId"]
-    else:
-        if task["code"] == "ResourceNotFound":
+        raise Exception(f"Unsupported platform: {platform}")
+
+    route = f"gecko.v2.{branch}.revision.{revision}.{product}.{platform_name}"
+    index = taskcluster_config.get_service("index")
+    try:
+        return index.findTask(route)["taskId"]
+    except taskcluster.exceptions.TaskclusterRestFailure as e:
+        if e.status_code == 404:
             return None
-        else:
-            raise TaskclusterException("Unknown TaskCluster index error.")
+        raise
 
 
 def get_task_details(task_id):
-    r = requests.get(queue_base + "task/{}".format(task_id))
-    r.raise_for_status()
-    return r.json()
+    queue = taskcluster_config.get_service("queue")
+    return queue.task(task_id)
 
 
 def get_task_status(task_id):
-    r = requests.get(queue_base + "task/{}/status".format(task_id))
-    r.raise_for_status()
-    return r.json()
+    queue = taskcluster_config.get_service("queue")
+    return queue.status(task_id)
 
 
 def get_task_artifacts(task_id):
-    r = requests.get(queue_base + "task/{}/artifacts".format(task_id))
-    r.raise_for_status()
-    return r.json()["artifacts"]
+    queue = taskcluster_config.get_service("queue")
+    return queue.listLatestArtifacts(task_id)["artifacts"]
 
 
 def get_tasks_in_group(group_id):
-    list_url = queue_base + "task-group/{}/list".format(group_id)
-
-    r = requests.get(list_url, params={"limit": 200})
-    r.raise_for_status()
-    reply = r.json()
-    tasks = reply["tasks"]
-    while "continuationToken" in reply:
-        r = requests.get(
-            list_url,
-            params={"limit": 200, "continuationToken": reply["continuationToken"]},
-        )
-        r.raise_for_status()
-        reply = r.json()
-        tasks += reply["tasks"]
-    return tasks
+    queue = taskcluster_config.get_service("queue")
+
+    token = None
+    while True:
+        query = {"limit": 200}
+        if token is not None:
+            query["continuationToken"] = token
+
+        response = queue.listTaskGroup(group_id, query=query)
+
+        for task in response["tasks"]:
+            yield task
+
+        token = response.get("continuationToken")
+        if token is None:
+            break
 
 
 def download_artifact(artifact_path, task_id, artifact_name):
     if os.path.exists(artifact_path):
         return artifact_path
 
-    def perform_download():
-        r = requests.get(
-            queue_base + "task/{}/artifacts/{}".format(task_id, artifact_name),
-            stream=True,
-        )
-
-        r.raise_for_status()
+    queue = taskcluster_config.get_service("queue")
 
+    def perform_download():
+        response = queue.getLatestArtifact(task_id, artifact_name)
         with open(artifact_path, "wb") as f:
-            r.raw.decode_content = True
-            shutil.copyfileobj(r.raw, f)
+            content = response["response"].content
+            f.write(content)
 
         if artifact_path.endswith(".zip") and not is_zipfile(artifact_path):
             raise BadZipFile("File is not a zip file")

bot/tests/test_artifacts.py

@@ -260,18 +260,23 @@ def build_task(task_state):
 
 
 def test_download_all(
-    LINUX_TASK_ID, LINUX_TASK, GROUP_TASKS_1, GROUP_TASKS_2, fake_artifacts
+    LINUX_TASK_ID,
+    LINUX_TASK,
+    GROUP_TASKS_1,
+    GROUP_TASKS_2,
+    fake_artifacts,
+    mock_taskcluster,
 ):
     responses.add(
         responses.GET,
-        f"https://queue.taskcluster.net/v1/task/{LINUX_TASK_ID}",
+        f"http://taskcluster.test/api/queue/v1/task/{LINUX_TASK_ID}",
         json=LINUX_TASK,
         status=200,
     )
     for group_tasks in _group_tasks():
         responses.add(
             responses.GET,
-            "https://queue.taskcluster.net/v1/task-group/aPt9FbIdQwmhwDIPDYLuaw/list",
+            "http://taskcluster.test/api/queue/v1/task-group/aPt9FbIdQwmhwDIPDYLuaw/list",
             json=group_tasks,
             status=200,
         )

bot/tests/test_taskcluster.py

@@ -6,37 +6,39 @@
 from zipfile import BadZipFile
 
 import pytest
-import requests
 import responses
+from taskcluster.exceptions import TaskclusterRestFailure
 
 from code_coverage_bot import taskcluster
 from conftest import FIXTURES_DIR
 
 
-def test_get_task_status(LINUX_TASK_ID, LINUX_TASK_STATUS):
+def test_get_task_status(mock_taskcluster, LINUX_TASK_ID, LINUX_TASK_STATUS):
     responses.add(
         responses.GET,
-        f"https://queue.taskcluster.net/v1/task/{LINUX_TASK_ID}/status",
+        f"http://taskcluster.test/api/queue/v1/task/{LINUX_TASK_ID}/status",
         json=LINUX_TASK_STATUS,
         status=200,
     )
     assert taskcluster.get_task_status(LINUX_TASK_ID) == LINUX_TASK_STATUS
 
 
-def test_get_task_details(LINUX_TASK_ID, LINUX_TASK):
+def test_get_task_details(mock_taskcluster, LINUX_TASK_ID, LINUX_TASK):
     responses.add(
         responses.GET,
-        f"https://queue.taskcluster.net/v1/task/{LINUX_TASK_ID}",
+        f"http://taskcluster.test/api/queue/v1/task/{LINUX_TASK_ID}",
         json=LINUX_TASK,
         status=200,
     )
     assert taskcluster.get_task_details(LINUX_TASK_ID) == LINUX_TASK
 
 
-def test_get_task(LINUX_TASK_ID, LATEST_LINUX, WIN_TASK_ID, LATEST_WIN):
+def test_get_task(
+    mock_taskcluster, LINUX_TASK_ID, LATEST_LINUX, WIN_TASK_ID, LATEST_WIN
+):
     responses.add(
         responses.GET,
-        "https://index.taskcluster.net/v1/task/gecko.v2.mozilla-central.revision.b2a9a4bb5c94de179ae7a3f52fde58c0e2897498.firefox.linux64-ccov-debug",
+        "http://taskcluster.test/api/index/v1/task/gecko.v2.mozilla-central.revision.b2a9a4bb5c94de179ae7a3f52fde58c0e2897498.firefox.linux64-ccov-debug",
         json=LATEST_LINUX,
         status=200,
     )  # noqa
@@ -49,7 +51,7 @@ def test_get_task(LINUX_TASK_ID, LATEST_LINUX, WIN_TASK_ID, LATEST_WIN):
 
     responses.add(
         responses.GET,
-        "https://index.taskcluster.net/v1/task/gecko.v2.mozilla-central.revision.916103b8675d9fdb28b891cac235d74f9f475942.firefox.win64-ccov-debug",
+        "http://taskcluster.test/api/index/v1/task/gecko.v2.mozilla-central.revision.916103b8675d9fdb28b891cac235d74f9f475942.firefox.win64-ccov-debug",
         json=LATEST_WIN,
         status=200,
     )  # noqa
@@ -61,10 +63,10 @@ def test_get_task(LINUX_TASK_ID, LATEST_LINUX, WIN_TASK_ID, LATEST_WIN):
     )
 
 
-def test_get_task_not_found(TASK_NOT_FOUND):
+def test_get_task_not_found(mock_taskcluster, TASK_NOT_FOUND):
     responses.add(
         responses.GET,
-        "https://index.taskcluster.net/v1/task/gecko.v2.mozilla-central.revision.b2a9a4bb5c94de179ae7a3f52fde58c0e2897498.firefox.linux64-ccov-debug",
+        "http://taskcluster.test/api/index/v1/task/gecko.v2.mozilla-central.revision.b2a9a4bb5c94de179ae7a3f52fde58c0e2897498.firefox.linux64-ccov-debug",
         json=TASK_NOT_FOUND,
         status=404,
     )  # noqa
@@ -77,28 +79,26 @@ def test_get_task_not_found(TASK_NOT_FOUND):
     )
 
 
-def test_get_task_failure(TASK_NOT_FOUND):
+def test_get_task_failure(mock_taskcluster, TASK_NOT_FOUND):
     err = TASK_NOT_FOUND.copy()
     err["code"] = "RandomError"
     responses.add(
         responses.GET,
-        "https://index.taskcluster.net/v1/task/gecko.v2.mozilla-central.revision.b2a9a4bb5c94de179ae7a3f52fde58c0e2897498.firefox.linux64-ccov-debug",
+        "http://taskcluster.test/api/index/v1/task/gecko.v2.mozilla-central.revision.b2a9a4bb5c94de179ae7a3f52fde58c0e2897498.firefox.linux64-ccov-debug",
         json=err,
         status=500,
     )  # noqa
 
-    with pytest.raises(
-        taskcluster.TaskclusterException, match="Unknown TaskCluster index error."
-    ):
+    with pytest.raises(TaskclusterRestFailure, match="Indexed task not found"):
         taskcluster.get_task(
             "mozilla-central", "b2a9a4bb5c94de179ae7a3f52fde58c0e2897498", "linux"
         )
 
 
-def test_get_task_artifacts(LINUX_TASK_ID, LINUX_TASK_ARTIFACTS):
+def test_get_task_artifacts(mock_taskcluster, LINUX_TASK_ID, LINUX_TASK_ARTIFACTS):
     responses.add(
         responses.GET,
-        f"https://queue.taskcluster.net/v1/task/{LINUX_TASK_ID}/artifacts",
+        f"http://taskcluster.test/api/queue/v1/task/{LINUX_TASK_ID}/artifacts",
         json=LINUX_TASK_ARTIFACTS,
         status=200,
     )
@@ -108,24 +108,24 @@ def test_get_task_artifacts(LINUX_TASK_ID, LINUX_TASK_ARTIFACTS):
     )
 
 
-def test_get_tasks_in_group(GROUP_TASKS_1, GROUP_TASKS_2):
+def test_get_tasks_in_group(mock_taskcluster, GROUP_TASKS_1, GROUP_TASKS_2):
     responses.add(
         responses.GET,
-        "https://queue.taskcluster.net/v1/task-group/aPt9FbIdQwmhwDIPDYLuaw/list?limit=200",
+        "http://taskcluster.test/api/queue/v1/task-group/aPt9FbIdQwmhwDIPDYLuaw/list?limit=200",
         json=GROUP_TASKS_1,
         status=200,
         match_querystring=True,
     )  # noqa
     responses.add(
         responses.GET,
-        "https://queue.taskcluster.net/v1/task-group/aPt9FbIdQwmhwDIPDYLuaw/list?continuationToken=1%2132%21YVB0OUZiSWRRd21od0RJUERZTHVhdw--~1%2132%21ZnJVcGRRT0VTalN0Nm9Ua1Ztcy04UQ--&limit=200",  # noqa
+        "http://taskcluster.test/api/queue/v1/task-group/aPt9FbIdQwmhwDIPDYLuaw/list?continuationToken=1%2132%21YVB0OUZiSWRRd21od0RJUERZTHVhdw--~1%2132%21ZnJVcGRRT0VTalN0Nm9Ua1Ztcy04UQ--&limit=200",  # noqa
         json=GROUP_TASKS_2,
         status=200,
         match_querystring=True,
     )  # noqa
 
     assert (
-        taskcluster.get_tasks_in_group("aPt9FbIdQwmhwDIPDYLuaw")
+        list(taskcluster.get_tasks_in_group("aPt9FbIdQwmhwDIPDYLuaw"))
         == GROUP_TASKS_1["tasks"] + GROUP_TASKS_2["tasks"]
     )
 
@@ -256,18 +256,15 @@ def test_get_platform(task_name, expected):
 
 
 @mock.patch("time.sleep")
-def test_download_artifact_forbidden(mocked_sleep, tmpdir):
+def test_download_artifact_forbidden(mocked_sleep, mock_taskcluster, tmpdir):
     responses.add(
         responses.GET,
-        "https://queue.taskcluster.net/v1/task/FBdocjnAQOW_GJDOfmgjxw/artifacts/public/test_info/code-coverage-grcov.zip",  # noqa
+        "http://taskcluster.test/api/queue/v1/task/FBdocjnAQOW_GJDOfmgjxw/artifacts/public%2Ftest_info%2Fcode-coverage-grcov.zip",  # noqa
         body="xml error...",
         status=403,
     )
 
-    with pytest.raises(
-        requests.exceptions.HTTPError,
-        match="403 Client Error: Forbidden for url: https://queue.taskcluster.net/v1/task/FBdocjnAQOW_GJDOfmgjxw/artifacts/public/test_info/code-coverage-grcov.zip",  # noqa
-    ):
+    with pytest.raises(TaskclusterRestFailure):
         taskcluster.download_artifact(
             os.path.join(tmpdir.strpath, "windows_reftest-6_code-coverage-grcov.zip"),
             "FBdocjnAQOW_GJDOfmgjxw",
@@ -278,10 +275,10 @@ def test_download_artifact_forbidden(mocked_sleep, tmpdir):
 
 
 @mock.patch("time.sleep")
-def test_download_artifact_badzip(mocked_sleep, tmpdir):
+def test_download_artifact_badzip(mocked_sleep, mock_taskcluster, tmpdir):
     responses.add(
         responses.GET,
-        "https://queue.taskcluster.net/v1/task/FBdocjnAQOW_GJDOfmgjxw/artifacts/public/test_info/code-coverage-grcov.zip",  # noqa
+        "http://taskcluster.test/api/queue/v1/task/FBdocjnAQOW_GJDOfmgjxw/artifacts/public%2Ftest_info%2Fcode-coverage-grcov.zip",  # noqa
         body="NOT A ZIP FILE",
         status=200,
         stream=True,