diff --git a/pyproject.toml b/pyproject.toml index f81ca71..1482506 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "bartender" -version = "0.5.2" +version = "0.5.3" description = "Job middleware for i-VRESSE" authors = [ diff --git a/src/bartender/web/api/job/views.py b/src/bartender/web/api/job/views.py index 9ff5cc5..0e87479 100644 --- a/src/bartender/web/api/job/views.py +++ b/src/bartender/web/api/job/views.py @@ -370,12 +370,13 @@ def _remove_archive(filename: str) -> None: }, response_class=FileResponse, ) -async def retrieve_job_directory_as_archive( +async def retrieve_job_directory_as_archive( # noqa: WPS211 job_dir: CurrentCompletedJobDir, background_tasks: BackgroundTasks, archive_format: ArchiveFormat = ".zip", exclude: Optional[list[str]] = Query(default=None), exclude_dirs: Optional[list[str]] = Query(default=None), + filename: Optional[str] = Query(default=None), # Note: also tried to with include (filter & filter_dirs) but that can be # unintuitive. e.g. include_dirs=['output'] doesn't return subdirs of # /output that are not also called output. Might improve when globs will be @@ -391,6 +392,8 @@ async def retrieve_job_directory_as_archive( '.tar.xz', '.tar.gz', '.tar.bz2' exclude: list of filename patterns that should be excluded from archive. exclude_dirs: list of directory patterns that should be excluded from archive. + filename: Name of the archive file to be returned. + If not provided, uses id of the job. Returns: FileResponse: Archive containing the content of job_dir @@ -402,6 +405,8 @@ async def retrieve_job_directory_as_archive( background_tasks.add_task(_remove_archive, archive_fn) return_fn = Path(archive_fn).name + if filename: + return_fn = filename return FileResponse(archive_fn, filename=return_fn) @@ -413,6 +418,7 @@ async def retrieve_job_subdirectory_as_archive( # noqa: WPS211 archive_format: ArchiveFormat = ".zip", exclude: Optional[list[str]] = Query(default=None), exclude_dirs: Optional[list[str]] = Query(default=None), + filename: Optional[str] = Query(default=None), ) -> FileResponse: """Download job output as archive. @@ -424,6 +430,8 @@ async def retrieve_job_subdirectory_as_archive( # noqa: WPS211 '.tar', '.tar.xz', '.tar.gz', '.tar.bz2' exclude: list of filename patterns that should be excluded from archive. exclude_dirs: list of directory patterns that should be excluded from archive. + filename: Name of the archive file to be returned. + If not provided, uses id of the job. Returns: FileResponse: Archive containing the output of job_dir @@ -436,6 +444,7 @@ async def retrieve_job_subdirectory_as_archive( # noqa: WPS211 archive_format=archive_format, exclude=exclude, exclude_dirs=exclude_dirs, + filename=filename, ) diff --git a/tests/web/test_job.py b/tests/web/test_job.py index 080c31a..bd59edc 100644 --- a/tests/web/test_job.py +++ b/tests/web/test_job.py @@ -758,6 +758,45 @@ async def test_job_directory_as_archive( assert stdout == "this is stdout" +@pytest.mark.anyio +@pytest.mark.parametrize( + "archive_format", + [".zip"], +) +async def test_job_directory_as_named_archive( + fastapi_app: FastAPI, + client: AsyncClient, + auth_headers: Dict[str, str], + mock_ok_job: int, + archive_format: str, +) -> None: + url = ( + fastapi_app.url_path_for( + "retrieve_job_directory_as_archive", + jobid=mock_ok_job, + ) + + f"?archive_format={archive_format}&filename=foo.zip" + ) + response = await client.get(url, headers=auth_headers) + + expected_content_type = ( + "application/zip" if archive_format == ".zip" else "application/x-tar" + ) + expected_content_disposition = f'attachment; filename="foo{archive_format}"' + + assert response.status_code == status.HTTP_200_OK + assert response.headers["content-type"] == expected_content_type + assert response.headers["content-disposition"] == expected_content_disposition + + fs = ZipFS if archive_format == ".zip" else TarFS + + with io.BytesIO(response.content) as responsefile: + with fs(responsefile) as archive: + stdout = archive.readtext("stdout.txt") + + assert stdout == "this is stdout" + + @pytest.mark.anyio async def test_job_subdirectory_as_archive( fastapi_app: FastAPI, @@ -785,6 +824,34 @@ async def test_job_subdirectory_as_archive( assert stdout == "hi from output dir" +@pytest.mark.anyio +async def test_job_subdirectory_as_named_archive( + fastapi_app: FastAPI, + client: AsyncClient, + auth_headers: Dict[str, str], + mock_ok_job: int, +) -> None: + url = ( + fastapi_app.url_path_for( + "retrieve_job_subdirectory_as_archive", + jobid=mock_ok_job, + path="output", + ) + + "?filename=bar.zip" + ) + response = await client.get(url, headers=auth_headers) + + assert response.status_code == status.HTTP_200_OK + assert response.headers["content-type"] == "application/zip" + assert response.headers["content-disposition"] == 'attachment; filename="bar.zip"' + + with io.BytesIO(response.content) as responsefile: + with ZipFS(responsefile) as archive: + stdout = archive.readtext("readme.txt") + + assert stdout == "hi from output dir" + + @pytest.fixture def demo_interactive_application( demo_config: Config,