diff --git a/src/openflexure_microscope_server/things/camera/__init__.py b/src/openflexure_microscope_server/things/camera/__init__.py index e541b11d..753aa66f 100644 --- a/src/openflexure_microscope_server/things/camera/__init__.py +++ b/src/openflexure_microscope_server/things/camera/__init__.py @@ -111,14 +111,19 @@ class CameraMemoryBuffer: try: image_tuple = list(self._storage.values())[-1] except IndexError as e: - raise NoImageInMemoryError("No image in memory to save.") from e + raise NoImageInMemoryError("No image in memory to retrieve.") from e # Clear the storage so images don't get retrieved out of order self._storage.clear() return image_tuple - if remove: - return self._storage.pop(buffer_id) - return self._storage[buffer_id] + try: + if remove: + return self._storage.pop(buffer_id) + return self._storage[buffer_id] + except KeyError as e: + raise NoImageInMemoryError( + "No image with matching id in memory to retrieve." + ) from e def clear(self): """ diff --git a/tests/test_camera_buffer.py b/tests/test_camera_buffer.py new file mode 100644 index 00000000..35c226f7 --- /dev/null +++ b/tests/test_camera_buffer.py @@ -0,0 +1,208 @@ +""" +Tests for the built in buffer for the camera. +""" + +from random import randint + +import numpy as np +from PIL import Image +import pytest + + +from openflexure_microscope_server.things.camera import ( + CameraMemoryBuffer, + NoImageInMemoryError, +) + +RANDOM_GENERATOR = np.random.default_rng() + + +def random_image(): + """Create a random image""" + imarray = RANDOM_GENERATOR.integers( + low=0, high=255, size=(100, 100, 3), dtype="uint8" + ) + return Image.fromarray(imarray) + + +def random_metadata(): + """Create a misc dictionary to pretend to be metadata""" + # Not very metadata like, but we are just the same dict it is returned + return {"a": randint(1, 100), "b": randint(1, 100)} + + +def test_add_and_get_image(): + """Check images can be captured and retrieved""" + mem_buf = CameraMemoryBuffer() + misc_image = random_image() + buffer_id = mem_buf.add_image(misc_image) + returned_image, _ = mem_buf.get_image(buffer_id) + # It is the same image + assert misc_image is returned_image + # It is now removed from memory + with pytest.raises(NoImageInMemoryError): + mem_buf.get_image(buffer_id) + + +def test_add_and_get_image_twice(): + """Check images can be retrieved twice if remove flag set false""" + mem_buf = CameraMemoryBuffer() + misc_image = random_image() + buffer_id = mem_buf.add_image(misc_image) + returned_image, _ = mem_buf.get_image(buffer_id, remove=False) + # It is the same image + assert misc_image is returned_image + # It is still in memory + returned_image, _ = mem_buf.get_image(buffer_id) + assert misc_image is returned_image + # It is now removed from memory + with pytest.raises(NoImageInMemoryError): + mem_buf.get_image(buffer_id) + + +def test_get_without_id(): + """Check images can be captured and retrieved without ID""" + mem_buf = CameraMemoryBuffer() + misc_image = random_image() + mem_buf.add_image(misc_image) + returned_image, _ = mem_buf.get_image() + # It is the same image + assert misc_image is returned_image + # It is now removed from memory + with pytest.raises(NoImageInMemoryError): + mem_buf.get_image() + + +def test_get_two_images(): + """Check two images can be retrieved.""" + mem_buf = CameraMemoryBuffer() + misc_image1 = random_image() + misc_image2 = random_image() + buffer_id1 = mem_buf.add_image(misc_image1, buffer_max=2) + buffer_id2 = mem_buf.add_image(misc_image2, buffer_max=2) + returned_image1, _ = mem_buf.get_image(buffer_id1) + returned_image2, _ = mem_buf.get_image(buffer_id2) + # It they the same images + assert misc_image1 is returned_image1 + assert misc_image2 is returned_image2 + # They are removed from memory + with pytest.raises(NoImageInMemoryError): + mem_buf.get_image(buffer_id1) + with pytest.raises(NoImageInMemoryError): + mem_buf.get_image(returned_image2) + + +def test_get_two_images_without_setting_buffer_size(): + """Check two images can't be retrieved if the buffer size isn't set""" + mem_buf = CameraMemoryBuffer() + misc_image1 = random_image() + misc_image2 = random_image() + buffer_id1 = mem_buf.add_image(misc_image1) + buffer_id2 = mem_buf.add_image(misc_image2) + with pytest.raises(NoImageInMemoryError): + mem_buf.get_image(buffer_id1) + returned_image2, _ = mem_buf.get_image(buffer_id2) + # Image 2 the expected image + assert misc_image2 is returned_image2 + + +def test_buffer_size_changing(): + """Check buffer size resets back to 1 when if not set""" + mem_buf = CameraMemoryBuffer() + misc_image1 = random_image() + misc_image2 = random_image() + misc_image3 = random_image() + buffer_id1 = mem_buf.add_image(misc_image1, buffer_max=3) + buffer_id2 = mem_buf.add_image(misc_image2, buffer_max=3) + # Third capture doen't set buffer size, so it will be reset + buffer_id3 = mem_buf.add_image(misc_image3) + # As buffer size was reset, images 1 and 2 are deleted + with pytest.raises(NoImageInMemoryError): + mem_buf.get_image(buffer_id1) + with pytest.raises(NoImageInMemoryError): + mem_buf.get_image(buffer_id2) + returned_image3, _ = mem_buf.get_image(buffer_id3) + # Image 2 the expected image + assert misc_image3 is returned_image3 + + +def test_capture_two_images_get_without_id(): + """Check that all images are deleted when getting without id""" + mem_buf = CameraMemoryBuffer() + misc_image1 = random_image() + misc_image2 = random_image() + mem_buf.add_image(misc_image1, buffer_max=2) + mem_buf.add_image(misc_image2, buffer_max=2) + returned_image, _ = mem_buf.get_image() + # It is the same image + assert returned_image is misc_image2 + # All were wiped from memory + with pytest.raises(NoImageInMemoryError): + mem_buf.get_image() + + +def test_buffer_size_respected(): + """Capture 10 images with a buffer size of 5. Check only last 5 exist""" + mem_buf = CameraMemoryBuffer() + + images = [] + buffer_ids = [] + for i in range(10): + image = random_image() + buffer_id = mem_buf.add_image(image, buffer_max=5) + images.append(image) + buffer_ids.append(buffer_id) + + for i, (image, buffer_id) in enumerate(zip(images, buffer_ids)): + if i < 5: + with pytest.raises(NoImageInMemoryError): + mem_buf.get_image(buffer_id) + else: + returned_image, _ = mem_buf.get_image(buffer_id) + assert image is returned_image + + +def test_clear_buffer(): + """Capture 10 images clear the buffer and check they are gone""" + mem_buf = CameraMemoryBuffer() + + images = [] + buffer_ids = [] + for i in range(10): + image = random_image() + buffer_id = mem_buf.add_image(image, buffer_max=10) + images.append(image) + buffer_ids.append(buffer_id) + + # Clear the buffer + mem_buf.clear() + + # They are now gone + for i, (image, buffer_id) in enumerate(zip(images, buffer_ids)): + with pytest.raises(NoImageInMemoryError): + mem_buf.get_image(buffer_id) + + +def test_get_metadata_too(): + """Capture 10 images clear the buffer and check they are gone""" + mem_buf = CameraMemoryBuffer() + + images = [] + metadatas = [] + buffer_ids = [] + for i in range(10): + image = random_image() + metadata = random_metadata() + buffer_id = mem_buf.add_image(image, metadata, buffer_max=10) + images.append(image) + metadatas.append(metadata) + buffer_ids.append(buffer_id) + + # Preallocate zipped data, to avoid long confusing lines + zipped = zip(images, metadatas, buffer_ids) + + # Check both image and metdata + for i, (image, metadata, buffer_id) in enumerate(zipped): + returned_image, returned_metadata = mem_buf.get_image(buffer_id) + assert image is returned_image + assert metadata is returned_metadata