From 944b748cedbe570558c2d8948ae766acf5be3373 Mon Sep 17 00:00:00 2001
From: y0z <y0z@users.noreply.github.com>
Date: Mon, 15 Jan 2024 18:26:00 +0900
Subject: [PATCH] Enhance delete artifacts APIs and add tests for them

---
 optuna_dashboard/artifact/_backend.py | 13 +++++--
 python_tests/artifact/test_backend.py | 49 +++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/optuna_dashboard/artifact/_backend.py b/optuna_dashboard/artifact/_backend.py
index 81cd766d..b8a44abd 100644
--- a/optuna_dashboard/artifact/_backend.py
+++ b/optuna_dashboard/artifact/_backend.py
@@ -14,6 +14,7 @@ from bottle import HTTPResponse
 from bottle import request
 from bottle import response
 import optuna
+from optuna.artifacts.exceptions import ArtifactNotFound
 from optuna.trial import FrozenTrial
 
 from .._bottle_util import json_api_view
@@ -184,7 +185,11 @@ def register_artifact_route(
         if artifact_store is None:
             response.status = 400  # Bad Request
             return {"reason": "Cannot access to the artifacts."}
-        artifact_store.remove(artifact_id)
+        try:
+            artifact_store.remove(artifact_id)
+        except ArtifactNotFound as e:
+            response.status = 404
+            return {"reason": str(e)}
 
         # The artifact's metadata is stored in one of the following two locations:
         storage.set_study_system_attr(
@@ -203,7 +208,11 @@ def register_artifact_route(
         if artifact_store is None:
             response.status = 400  # Bad Request
             return {"reason": "Cannot access to the artifacts."}
-        artifact_store.remove(artifact_id)
+        try:
+            artifact_store.remove(artifact_id)
+        except ArtifactNotFound as e:
+            response.status = 404
+            return {"reason": str(e)}
 
         storage.set_study_system_attr(
             study_id, ARTIFACTS_ATTR_PREFIX + artifact_id, json.dumps(None)
diff --git a/python_tests/artifact/test_backend.py b/python_tests/artifact/test_backend.py
index 52fe6027..165c5029 100644
--- a/python_tests/artifact/test_backend.py
+++ b/python_tests/artifact/test_backend.py
@@ -249,3 +249,52 @@ def test_upload_artifact() -> None:
         with open(f"{tmpdir}/{res['artifact_id']}", "r") as f:
             data = f.read()
             assert data == "dummy_content"
+
+
+def test_delete_study_artifact() -> None:
+    storage = optuna.storages.InMemoryStorage()
+    study = optuna.create_study(storage=storage)
+    with tempfile.TemporaryDirectory() as tmpdir:
+        artifact_store = FileSystemArtifactStore(tmpdir)
+        with tempfile.NamedTemporaryFile() as f:
+            f.write(b"dummy_content")
+            f.flush()
+            artifact_id = upload_artifact(study, f.name, artifact_store=artifact_store)
+        app = create_app(storage, artifact_store)
+        status, _, _ = send_request(
+            app,
+            f"/api/artifacts/{study._study_id}/{artifact_id}",
+            "DELETE",
+        )
+        assert status == 204
+        status, _, _ = send_request(
+            app,
+            f"/api/artifacts/{study._study_id}/{artifact_id}",
+            "DELETE",
+        )
+        assert status == 404
+
+
+def test_delete_trial_artifact() -> None:
+    storage = optuna.storages.InMemoryStorage()
+    study = optuna.create_study(storage=storage)
+    trial = study.ask()
+    with tempfile.TemporaryDirectory() as tmpdir:
+        artifact_store = FileSystemArtifactStore(tmpdir)
+        with tempfile.NamedTemporaryFile() as f:
+            f.write(b"dummy_content")
+            f.flush()
+            artifact_id = upload_artifact(trial, f.name, artifact_store=artifact_store)
+        app = create_app(storage, artifact_store)
+        status, _, _ = send_request(
+            app,
+            f"/api/artifacts/{study._study_id}/{trial._trial_id}/{artifact_id}",
+            "DELETE",
+        )
+        assert status == 204
+        status, _, _ = send_request(
+            app,
+            f"/api/artifacts/{study._study_id}/{trial._trial_id}/{artifact_id}",
+            "DELETE",
+        )
+        assert status == 404