From 5c295da6ffe758730662918f639dc296b9d0a96a Mon Sep 17 00:00:00 2001
From: Sebastian Schmidt <mrschmidt@google.com>
Date: Tue, 28 Mar 2023 10:24:16 -0700
Subject: [PATCH] Return custom error if model download fails

PiperOrigin-RevId: 520066065
---
 mediapipe/tasks/web/core/task_runner.ts      |  9 ++++++++-
 mediapipe/tasks/web/core/task_runner_test.ts | 12 +++++++++++-
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/mediapipe/tasks/web/core/task_runner.ts b/mediapipe/tasks/web/core/task_runner.ts
index 8f69423fb..238564f1d 100644
--- a/mediapipe/tasks/web/core/task_runner.ts
+++ b/mediapipe/tasks/web/core/task_runner.ts
@@ -114,7 +114,14 @@ export abstract class TaskRunner {
       // We don't use `await` here since we want to apply most settings
       // synchronously.
       return fetch(baseOptions.modelAssetPath.toString())
-          .then(response => response.arrayBuffer())
+          .then(response => {
+            if (!response.ok) {
+              throw new Error(`Failed to fetch model: ${
+                  baseOptions.modelAssetPath} (${response.status})`);
+            } else {
+              return response.arrayBuffer();
+            }
+          })
           .then(buffer => {
             this.setExternalFile(new Uint8Array(buffer));
             this.refreshGraph();
diff --git a/mediapipe/tasks/web/core/task_runner_test.ts b/mediapipe/tasks/web/core/task_runner_test.ts
index 41ac76952..f44cfe49e 100644
--- a/mediapipe/tasks/web/core/task_runner_test.ts
+++ b/mediapipe/tasks/web/core/task_runner_test.ts
@@ -118,12 +118,14 @@ describe('TaskRunner', () => {
 
   let fetchSpy: jasmine.Spy;
   let taskRunner: TaskRunnerFake;
+  let fetchStatus = 200;
 
   beforeEach(() => {
     fetchSpy = jasmine.createSpy().and.callFake(async url => {
-      expect(url).toEqual('foo');
       return {
         arrayBuffer: () => mockBytes.buffer,
+        ok: fetchStatus === 200,
+        status: fetchStatus,
       } as unknown as Response;
     });
     global.fetch = fetchSpy;
@@ -225,6 +227,14 @@ describe('TaskRunner', () => {
     return resolvedPromise;
   });
 
+  it('returns custom error if model download failed', () => {
+    fetchStatus = 404;
+    return expectAsync(taskRunner.setOptions({
+             baseOptions: {modelAssetPath: `notfound.tflite`}
+           }))
+        .toBeRejectedWithError('Failed to fetch model: notfound.tflite (404)');
+  });
+
   it('can enable CPU delegate', async () => {
     await taskRunner.setOptions({
       baseOptions: {