From 4666f787a0e06e55febde3723a5c9a586fb32856 Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@pm.me>
Date: Mon, 23 Mar 2020 15:29:01 -0400
Subject: [PATCH] Raise an 'ArtifactNotFound' exception when requesting a
 non-existent artifact

---
 mozci/errors.py           | 18 +++++++++++++++++-
 mozci/task.py             | 14 +++++++++++++-
 mozci/util/taskcluster.py |  6 +++++-
 tests/test_task.py        | 40 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 75 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_task.py

diff --git a/mozci/errors.py b/mozci/errors.py
index d81cc2f8..28c8a596 100644
--- a/mozci/errors.py
+++ b/mozci/errors.py
@@ -5,7 +5,7 @@ class BasePushException(Exception):
     def __init__(self, rev, branch, msg):
         self.rev = rev
         self.branch = branch
-        self.msg = f"Error processing push '{rev}' on {branch}: {msg}"
+        self.msg = f"Error with push '{rev}' on {branch}: {msg}"
 
 
 class PushNotFound(BasePushException):
@@ -14,3 +14,19 @@ class PushNotFound(BasePushException):
     def __init__(self, *args, **kwargs):
         kwargs["msg"] = "does not exist!"
         super(PushNotFound, self).__init__(*args, **kwargs)
+
+
+class BaseTaskException(Exception):
+    def __init__(self, id, label, msg):
+        self.id = id
+        self.label = label
+        self.msg = f"Error with task '{id}' ({label}): {msg}"
+
+
+class ArtifactNotFound(BaseTaskException):
+    """Raised when the requested task artifact does not exist."""
+
+    def __init__(self, artifact, *args, **kwargs):
+        kwargs["msg"] = f"artifact '{artifact}' does not exist!"
+        self.artifact = artifact
+        super(ArtifactNotFound, self).__init__(*args, **kwargs)
diff --git a/mozci/task.py b/mozci/task.py
index a168b7be..2731f829 100644
--- a/mozci/task.py
+++ b/mozci/task.py
@@ -6,10 +6,12 @@
 from enum import Enum
 from typing import Dict, List
 
+import requests
 from adr.util import memoize, memoized_property
 from loguru import logger
 from urllib3.response import HTTPResponse
 
+from mozci.errors import ArtifactNotFound
 from mozci.util.taskcluster import get_artifact, list_artifacts
 
 
@@ -94,8 +96,18 @@ def get_artifact(self, path):
 
         Returns:
             Contents of the artifact.
+
+        Raises:
+            mozci.errors.ArtifactNotFound: When the requested artifact does not
+                                           exist.
         """
-        data = get_artifact(self.id, path)
+        try:
+            data = get_artifact(self.id, path)
+        except requests.exceptions.HTTPError as e:
+            if e.response.status_code == 404:
+                raise ArtifactNotFound(path, self.id, self.label) from e
+            raise
+
         if not isinstance(data, HTTPResponse):
             return data
 
diff --git a/mozci/util/taskcluster.py b/mozci/util/taskcluster.py
index 2505f238..085fe25f 100644
--- a/mozci/util/taskcluster.py
+++ b/mozci/util/taskcluster.py
@@ -60,8 +60,12 @@ def get_artifact(task_id, path):
     """
     try:
         response = _do_request(get_artifact_url(task_id, path))
-    except requests.exceptions.HTTPError:
+    except requests.exceptions.HTTPError as e:
+        if e.response.status_code != 404:
+            raise
+
         response = _do_request(get_artifact_url(task_id, path, old_deployment=True))
+
     return _handle_artifact(path, response)
 
 
diff --git a/tests/test_task.py b/tests/test_task.py
new file mode 100644
index 00000000..87b878a7
--- /dev/null
+++ b/tests/test_task.py
@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+
+import pytest
+
+from mozci.errors import ArtifactNotFound
+from mozci.task import Task
+from mozci.util.taskcluster import get_artifact_url
+
+
+@pytest.fixture
+def create_task():
+    id = 0
+
+    def inner(**kwargs):
+        nonlocal id
+        task = Task.create(id=id, **kwargs)
+        id += 1
+        return task
+
+    return inner
+
+
+def test_missing_artifacts(responses, create_task):
+    artifact = "public/artifact.txt"
+    task = create_task(label="foobar")
+
+    # First we'll check the new deployment.
+    responses.add(
+        responses.GET, get_artifact_url(task.id, artifact), status=404,
+    )
+
+    # Then we'll check the old deployment.
+    responses.add(
+        responses.GET,
+        get_artifact_url(task.id, artifact, old_deployment=True),
+        status=404,
+    )
+
+    with pytest.raises(ArtifactNotFound):
+        task.get_artifact(artifact)