diff --git a/python_tests/artifact/test_backend.py b/python_tests/artifact/test_backend.py index 22544e6f..00e55d51 100644 --- a/python_tests/artifact/test_backend.py +++ b/python_tests/artifact/test_backend.py @@ -1,9 +1,16 @@ +import tempfile from unittest.mock import MagicMock +import optuna +from optuna.artifacts import FileSystemArtifactStore +from optuna.artifacts import upload_artifact from optuna.storages import BaseStorage +from optuna_dashboard._app import create_app from optuna_dashboard.artifact import _backend import pytest +from ..wsgi_client import send_request + def test_get_artifact_path() -> None: study = MagicMock(_study_id=0) @@ -80,3 +87,50 @@ def test_list_trial_artifacts(init_storage_with_artifact_meta: MagicMock) -> Non {"artifact_id": "id1", "filename": "bar.txt"}, {"artifact_id": "id2", "filename": "baz.txt"}, ] + + +def test_artifact_store_none() -> None: + storage = optuna.storages.InMemoryStorage() + app = create_app(storage) + status, _, body = send_request( + app, + "/artifacts/0/0", + "GET", + content_type="application/json", + ) + assert status == 400 + + +def test_artifact_not_found() -> None: + storage = optuna.storages.InMemoryStorage() + study = optuna.create_study(storage=storage) + with tempfile.TemporaryDirectory() as tmpdir: + artifact_store = FileSystemArtifactStore(tmpdir) + app = create_app(storage, artifact_store) + status, _, body = send_request( + app, + f"/artifacts/{study._study_id}/abc123", + "GET", + content_type="application/json", + ) + assert status == 404 + + +def test_successful_artifact_retrieval() -> 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, _, body = send_request( + app, + f"/artifacts/{study._study_id}/{artifact_id}", + "GET", + content_type="application/json", + ) + assert status == 200 + assert body == b"dummy_content"