From 6b86fa0d29c9a0678c6802057128ded7c893da4f Mon Sep 17 00:00:00 2001
From: rahul-flex <rahul@flexcompute.com>
Date: Mon, 31 Mar 2025 23:19:18 -0400
Subject: [PATCH] Fix:Error message for invalid task ID

---
 CHANGELOG.md                  |  3 +++
 tests/test_web/test_webapi.py | 16 ++++++++++++++++
 tidy3d/web/core/exceptions.py |  6 ++++++
 tidy3d/web/core/http_util.py  |  4 ++--
 tidy3d/web/core/task_core.py  | 11 +++++++++--
 5 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3fafb001db..cd3292d3b0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `fill` and `fill_structures` argument in `td.Simulation.plot_structures()` and `td.Simulation.plot()` respectively to disable fill and plot outlines of structures only.
 - New subpixel averaging option `ContourPathAveraging` applied to dielectric material boundaries.
 
+### Changed
+- Error message for invalid task ID.
+
 ### Fixed
 - Compatibility with `xarray>=2025.03`.
 
diff --git a/tests/test_web/test_webapi.py b/tests/test_web/test_webapi.py
index 527915b05d..eeb4c542cb 100644
--- a/tests/test_web/test_webapi.py
+++ b/tests/test_web/test_webapi.py
@@ -39,6 +39,7 @@
     upload,
 )
 from tidy3d.web.core.environment import Env
+from tidy3d.web.core.exceptions import WebNotFoundError
 from tidy3d.web.core.types import TaskType
 
 TASK_NAME = "task_name_test"
@@ -722,3 +723,18 @@ def save_sim_to_path(path: str) -> None:
                 "--inspect_sim",
             ]
         )
+
+
+@responses.activate
+def test_load_invalid_task_raises(mock_webapi):
+    """Ensure that load() raises TaskNotFoundError for a non-existent task ID."""
+    fake_id = "INVALID_TASK_ID"
+
+    responses.add(
+        responses.GET,
+        f"{Env.current.web_api_endpoint}/tidy3d/tasks/{fake_id}/detail",
+        json={"error": "Task not found"},
+        status=404,
+    )
+    with pytest.raises(WebNotFoundError, match="Resource not found"):
+        load(fake_id)
diff --git a/tidy3d/web/core/exceptions.py b/tidy3d/web/core/exceptions.py
index 6eefaec4a0..398550eca1 100644
--- a/tidy3d/web/core/exceptions.py
+++ b/tidy3d/web/core/exceptions.py
@@ -11,3 +11,9 @@ def __init__(self, message: str = None):
         log = get_logger()
         super().__init__(message)
         log.error(message)
+
+
+class WebNotFoundError(WebError):
+    """A generic error indicating an HTTP 404 (resource not found)."""
+
+    pass
diff --git a/tidy3d/web/core/http_util.py b/tidy3d/web/core/http_util.py
index d30ec0f6d2..1535959465 100644
--- a/tidy3d/web/core/http_util.py
+++ b/tidy3d/web/core/http_util.py
@@ -25,7 +25,7 @@
 )
 from .core_config import get_logger
 from .environment import Env
-from .exceptions import WebError
+from .exceptions import WebError, WebNotFoundError
 
 REINITIALIZED = False
 
@@ -131,7 +131,7 @@ def wrapper(*args, **kwargs):
 
         if resp.status_code != ResponseCodes.OK.value:
             if resp.status_code == ResponseCodes.NOT_FOUND.value:
-                return None
+                raise WebNotFoundError("Resource not found (HTTP 404).")
             json_resp = resp.json()
             if "error" in json_resp.keys():
                 raise WebError(json_resp["error"])
diff --git a/tidy3d/web/core/task_core.py b/tidy3d/web/core/task_core.py
index beb50598d1..bd036d9d32 100644
--- a/tidy3d/web/core/task_core.py
+++ b/tidy3d/web/core/task_core.py
@@ -12,12 +12,14 @@
 from botocore.exceptions import ClientError
 from pydantic.v1 import Extra, Field, parse_obj_as
 
+import tidy3d as td
+
 from . import http_util
 from .cache import FOLDER_CACHE
 from .constants import SIM_ERROR_FILE, SIM_FILE_HDF5_GZ, SIM_LOG_FILE, SIMULATION_DATA_HDF5_GZ
 from .core_config import get_logger_console
 from .environment import Env
-from .exceptions import WebError
+from .exceptions import WebError, WebNotFoundError
 from .file_util import read_simulation_from_hdf5
 from .http_util import http
 from .s3utils import download_file, download_gz_file, upload_file
@@ -261,7 +263,12 @@ def get(cls, task_id: str, verbose: bool = True) -> SimulationTask:
             :class:`.SimulationTask` object containing info about status,
              size, credits of task and others.
         """
-        resp = http.get(f"tidy3d/tasks/{task_id}/detail")
+        try:
+            resp = http.get(f"tidy3d/tasks/{task_id}/detail")
+        except WebNotFoundError as e:
+            td.log.error(f"The requested task ID '{task_id}' does not exist.")
+            raise e
+
         task = SimulationTask(**resp) if resp else None
         return task