Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions plotting-service/plotting_service/plotting_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
from fastapi import FastAPI, HTTPException
from h5grove.fastapi_utils import router, settings # type: ignore
from starlette.middleware.cors import CORSMiddleware
from starlette.middleware.gzip import GZipMiddleware
from starlette.requests import Request

from plotting_service.auth import get_experiments_for_user, get_user_from_token
Expand Down Expand Up @@ -48,6 +49,7 @@
allow_methods=["*"],
allow_headers=["*"],
)
app.add_middleware(GZipMiddleware, minimum_size=1000)

CEPH_DIR = os.environ.get("CEPH_DIR", "/ceph")
logger.info("Setting ceph directory to %s", CEPH_DIR)
Expand Down
86 changes: 85 additions & 1 deletion plotting-service/plotting_service/routers/imat.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,16 +7,20 @@
from pathlib import Path

from fastapi import APIRouter, HTTPException, Query
from starlette.responses import JSONResponse
from PIL import Image
from starlette.responses import JSONResponse, Response

from plotting_service.services.image_service import (
IMAGE_SUFFIXES,
convert_image_to_rgb_array,
find_latest_image_in_directory,
)
from plotting_service.utils import safe_check_filepath

ImatRouter = APIRouter()

IMAT_DIR: Path = Path(os.getenv("IMAT_DIR", "/imat")).resolve()
CEPH_DIR = os.environ.get("CEPH_DIR", "/ceph")

stdout_handler = logging.StreamHandler(stream=sys.stdout)
logging.basicConfig(
Expand Down Expand Up @@ -87,3 +91,83 @@ async def get_latest_imat_image(
"downsampleFactor": effective_downsample,
}
return JSONResponse(payload)


@ImatRouter.get("/imat/list-images", summary="List images in a directory")
async def list_imat_images(
path: typing.Annotated[
str, Query(..., description="Path to the directory containing images, relative to CEPH_DIR")
],
) -> list[str]:
"""Return a sorted list of TIFF images in the given directory."""

dir_path = (Path(CEPH_DIR) / path).resolve()
# Security: Ensure path is within CEPH_DIR
try:
safe_check_filepath(dir_path, CEPH_DIR)
except FileNotFoundError as err:
raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail="Directory not found") from err

if not dir_path.exists() or not dir_path.is_dir():
raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail="Directory not found")

images = [entry.name for entry in dir_path.iterdir() if entry.is_file() and entry.suffix.lower() in IMAGE_SUFFIXES]

return sorted(images)


@ImatRouter.get("/imat/image", summary="Fetch a specific TIFF image as raw data")
async def get_imat_image(
path: typing.Annotated[str, Query(..., description="Path to the TIFF image file, relative to CEPH_DIR")],
downsample_factor: typing.Annotated[
int,
Query(
ge=1,
le=64,
description="Integer factor to reduce each dimension by (1 keeps original resolution).",
),
] = 1,
) -> Response:
"""Return the raw data of a specific TIFF image as binary."""

image_path = (Path(CEPH_DIR) / path).resolve()
# Security: Ensure path is within CEPH_DIR
try:
safe_check_filepath(image_path, CEPH_DIR)
except FileNotFoundError as err:
raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail="Directory not found") from err

if not image_path.exists() or not image_path.is_file():
raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail="Image not found")

try:
with Image.open(image_path) as img:
original_width, original_height = img.size

if downsample_factor > 1:
target_width = max(1, round(original_width / downsample_factor))
target_height = max(1, round(original_height / downsample_factor))
display_img = img.resize((target_width, target_height), Image.Resampling.NEAREST)
else:
display_img = img

sampled_width, sampled_height = display_img.size
# For 16-bit TIFFs, tobytes() returns raw 16-bit bytes
data_bytes = display_img.tobytes()

headers = {
"X-Image-Width": str(sampled_width),
"X-Image-Height": str(sampled_height),
"X-Original-Width": str(original_width),
"X-Original-Height": str(original_height),
"X-Downsample-Factor": str(downsample_factor),
"Access-Control-Expose-Headers": (
"X-Image-Width, X-Image-Height, X-Original-Width, X-Original-Height, X-Downsample-Factor"
),
}

return Response(content=data_bytes, media_type="application/octet-stream", headers=headers)

except Exception as exc:
logger.error(f"Failed to process image {image_path}: {exc}")
raise HTTPException(HTTPStatus.INTERNAL_SERVER_ERROR, "Unable to process image") from exc
2 changes: 1 addition & 1 deletion plotting-service/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ dependencies = [

[project.optional-dependencies]
formatting = [
"ruff==0.12.3",
"ruff==0.15.0",
"mypy==1.17.0",
"plotting-service[test]",
"types-requests==2.32.4.20250611",
Expand Down
147 changes: 147 additions & 0 deletions plotting-service/test/test_plotting_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -204,3 +204,150 @@ def test_get_latest_imat_image_with_mock_rb_folder(tmp_path, monkeypatch):
assert payload["shape"] == [2, 4, 3]
assert payload["downsampleFactor"] == 2 # noqa: PLR2004
assert payload["data"] == expected_bytes


def test_list_imat_images(tmp_path, monkeypatch):
"""Verify that /imat/list-images correctly filters and sorts image files."""
monkeypatch.setattr(imat, "CEPH_DIR", str(tmp_path))

# Create test directory
data_dir = tmp_path / "test_data"
data_dir.mkdir()
(data_dir / "image2.tiff").touch()
(data_dir / "image1.tif").touch()
(data_dir / "not_an_image.txt").touch()

client = TestClient(plotting_api.app)
response = client.get("/imat/list-images", params={"path": "test_data"}, headers={"Authorization": "Bearer foo"})

assert response.status_code == HTTPStatus.OK
assert response.json() == ["image1.tif", "image2.tiff"]


def test_list_imat_images_not_found(tmp_path, monkeypatch):
"""Ensure 404 is returned when the requested directory does not exist."""
monkeypatch.setattr(imat, "CEPH_DIR", str(tmp_path))

client = TestClient(plotting_api.app)
response = client.get("/imat/list-images", params={"path": "non_existent"}, headers={"Authorization": "Bearer foo"})

assert response.status_code == HTTPStatus.NOT_FOUND


def test_list_imat_images_forbidden(tmp_path, monkeypatch):
"""Verify that path traversal attempts are blocked with 403 Forbidden."""
monkeypatch.setattr(imat, "CEPH_DIR", str(tmp_path))

client = TestClient(plotting_api.app)
# safe_check_filepath should block this
response = client.get("/imat/list-images", params={"path": "../.."}, headers={"Authorization": "Bearer foo"})

assert response.status_code == HTTPStatus.FORBIDDEN


def test_get_imat_image(tmp_path, monkeypatch):
"""Ensure /imat/image returns raw binary data and correct metadata headers."""
monkeypatch.setattr(imat, "CEPH_DIR", str(tmp_path))

# Create a tiny 16-bit TIFF (4x4, all 1000)
image_path = tmp_path / "test.tif"
image = Image.new("I;16", (4, 4), color=1000)
image.save(image_path, format="TIFF")
image.close()

client = TestClient(plotting_api.app)
response = client.get("/imat/image", params={"path": "test.tif"}, headers={"Authorization": "Bearer foo"})

assert response.status_code == HTTPStatus.OK
assert response.headers["X-Image-Width"] == "4"
assert response.headers["X-Image-Height"] == "4"
assert response.headers["X-Original-Width"] == "4"
assert response.headers["X-Original-Height"] == "4"
assert response.headers["X-Downsample-Factor"] == "1"
assert response.content == Image.open(image_path).tobytes()
assert response.headers["Content-Type"] == "application/octet-stream"


def test_get_imat_image_downsampled(tmp_path, monkeypatch):
"""Verify that downsampling works and headers reflect the sampled dimensions."""
monkeypatch.setattr(imat, "CEPH_DIR", str(tmp_path))

image_path = tmp_path / "test.tif"
image = Image.new("I;16", (8, 4), color=1000)
image.save(image_path, format="TIFF")
image.close()

client = TestClient(plotting_api.app)
response = client.get(
"/imat/image", params={"path": "test.tif", "downsample_factor": 2}, headers={"Authorization": "Bearer foo"}
)

assert response.status_code == HTTPStatus.OK
assert response.headers["X-Image-Width"] == "4"
assert response.headers["X-Image-Height"] == "2"
assert response.headers["X-Original-Width"] == "8"
assert response.headers["X-Original-Height"] == "4"


def test_get_latest_imat_image_no_rb_folders(tmp_path, monkeypatch):
"""Ensure 404 is returned if no RB folders are present in the IMAT directory."""
monkeypatch.setattr(imat, "IMAT_DIR", tmp_path)

client = TestClient(plotting_api.app)
response = client.get("/imat/latest-image", headers={"Authorization": "Bearer foo"})

assert response.status_code == HTTPStatus.NOT_FOUND
assert "No RB folders" in response.json()["detail"]


def test_get_imat_image_not_found(tmp_path, monkeypatch):
"""Ensure /imat/image returns 404 for a missing file."""
monkeypatch.setattr(imat, "CEPH_DIR", str(tmp_path))

client = TestClient(plotting_api.app)
response = client.get("/imat/image", params={"path": "not_there.tif"}, headers={"Authorization": "Bearer foo"})

assert response.status_code == HTTPStatus.NOT_FOUND


def test_get_latest_imat_image_no_images_in_rb(tmp_path, monkeypatch):
"""Ensure 404 is returned if RB folders exist but contain no valid image files."""
monkeypatch.setattr(imat, "IMAT_DIR", tmp_path)
(tmp_path / "RB1234").mkdir()

client = TestClient(plotting_api.app)
response = client.get("/imat/latest-image", headers={"Authorization": "Bearer foo"})
assert response.status_code == HTTPStatus.NOT_FOUND
assert "No images found" in response.json()["detail"]


def test_get_imat_image_internal_error(tmp_path, monkeypatch):
"""Verify that an exception during image processing returns a 500 error."""
monkeypatch.setattr(imat, "CEPH_DIR", str(tmp_path))
(tmp_path / "corrupt.tif").touch()

client = TestClient(plotting_api.app)
with mock.patch("PIL.Image.open", side_effect=Exception("Simulated failure")):
response = client.get("/imat/image", params={"path": "corrupt.tif"}, headers={"Authorization": "Bearer foo"})

assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
assert "Unable to process image" in response.json()["detail"]


def test_get_latest_imat_image_conversion_error(tmp_path, monkeypatch):
"""Verify that an error during RB latest image conversion returns a 500 error."""
monkeypatch.setattr(imat, "IMAT_DIR", tmp_path)
rb_dir = tmp_path / "RB1234"
rb_dir.mkdir()
(rb_dir / "test.tif").touch()

client = TestClient(plotting_api.app)
with mock.patch(
"plotting_service.routers.imat.convert_image_to_rgb_array", side_effect=Exception("Conversion failed")
):
response = client.get(
"/imat/latest-image", params={"downsample_factor": 1}, headers={"Authorization": "Bearer foo"}
)

assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
assert "Unable to convert IMAT image" in response.json()["detail"]