openflexure-microscope-server/tests/test_scan_directories.py

"""Test the functionality in the scan_directories module."""

import tempfile
import os
import math
import shutil
import logging
import random
import time
import json
from collections import namedtuple

import pytest

from openflexure_microscope_server.scan_directories import (
    ScanDirectoryManager,
    ScanDirectory,
    ScanInfo,
    get_files_in_zip,
    NotEnoughFreeSpaceError,
    SCAN_DATA_FILENAME,
)

from .test_scan_data import _fake_scan_data
from .utilities import assert_unique_of_length

# A global logger to pass in as an Invocation Logger
LOGGER = logging.getLogger("mock-invocation_logger")

# Use our own dir in the root temp dir not a dynamically generated one so we
# have some control of when it is deleted
BASE_SCAN_DIR = os.path.join(tempfile.gettempdir(), "scans")


def _clear_scan_dir() -> None:
    """Delete the scan dir."""
    if os.path.exists(BASE_SCAN_DIR):
        shutil.rmtree(BASE_SCAN_DIR)


def _add_fake_image(scan_dir: ScanDirectory) -> None:
    """Make a fake image on disk in the scan directory."""
    unique = False
    while not unique:
        x_pos = random.randint(-10000, 10000)
        y_pos = random.randint(-10000, 10000)
        filename = f"image_{x_pos}_{y_pos}.jpg"
        filepath = os.path.join(scan_dir.images_dir, filename)
        unique = not os.path.exists(filepath)
    with open(filepath, "w") as f_obj:
        f_obj.write("fake")


def _add_fake_file(
    scan_dir: ScanDirectory, filename: str, in_im_dir: bool = False
) -> None:
    """Make a fake file on disk in the scan directory.

    :param in_im_dir: Boolean, if set True the fake file is created in the images
    directory of the scan not the root directory.
    """
    if in_im_dir:
        filepath = os.path.join(scan_dir.images_dir, filename)
    else:
        filepath = os.path.join(scan_dir.dir_path, filename)
    with open(filepath, "w") as f_obj:
        f_obj.write("fake")


def _make_fake_dzi(scan_dir: ScanDirectory, n_layers: int = 8) -> None:
    """Create a fake DZI in a scan.

    :param n_layers: The number of layers of tiles. I.e. tile directories numbered
        0...(n_layers-1) will be created. Default 8
    """
    # Add an a dzi image
    dzi_fname = scan_dir.name + ".dzi"
    dzi_path = os.path.join(scan_dir.images_dir, dzi_fname)
    with open(dzi_path, "w") as f_obj:
        f_obj.write("This should be xml")

    # and the directory for the dzi tiles
    dzi_tile_dir = scan_dir.name + "_files"
    dzi_tile_dir_path = os.path.join(scan_dir.images_dir, dzi_tile_dir)
    os.makedirs(dzi_tile_dir_path)

    # this directory then contains a directories numbered from 0-n
    for i in range(n_layers):
        layer_dir_path = os.path.join(dzi_tile_dir_path, str(i))
        os.makedirs(layer_dir_path)

        # Number of images (in each axis) in this layer. Number of tiles doubles each
        # layer, but the first few layers always have 1 image.
        n_ims = math.ceil(2 ** (i - 4))
        for x_index in range(n_ims):
            for y_index in range(n_ims):
                tile_name = f"{x_index}_{y_index}.jpg"
                tile_path = os.path.join(layer_dir_path, tile_name)
                with open(tile_path, "w") as f_obj:
                    f_obj.write("This would normally be jpeg data")


def test_basic_directory_operations():
    """Test some basic operations.

    Test some basic operations, including:
    - ScanDirectoryManager creates a scan directory
    - For a fake (dummy) scan, a path can be created and retrieved
    - For a fake (dummy) file, a path can be created and retrieved (for both a .zip and .img file)
    - When a scan directory is created, the directory exists as does the image directory
    - Scan directories can be deleted and scans can be cleared
    """
    _clear_scan_dir()
    assert not os.path.isdir(BASE_SCAN_DIR)
    scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR)
    # Check it makes the scan directory
    assert os.path.isdir(BASE_SCAN_DIR)

    # Considering a fake scan
    scan_name = "fake_scan_0001"
    scan_path = os.path.join(BASE_SCAN_DIR, scan_name)
    scan_im_dir = os.path.join(BASE_SCAN_DIR, scan_name, "images")

    # It doesn't exist but we can check its paths
    assert not scan_dir_manager.exists(scan_name)
    assert not os.path.isdir(scan_path)
    assert scan_dir_manager.path_for(scan_name) == scan_path
    assert scan_dir_manager.img_dir_for(scan_name) == scan_im_dir

    # Get the path of a fake file
    fake_file = scan_dir_manager.get_file_path_from(scan_name, "foo.zip")
    assert fake_file == os.path.join(scan_path, "foo.zip")
    # But this is none if we check it exists
    fake_file = scan_dir_manager.get_file_path_from(
        scan_name, "foo.zip", check_exists=True
    )
    assert fake_file is None

    # Get the path of another fake file
    fake_file = scan_dir_manager.get_file_path_from_img_dir(scan_name, "bar.img")
    assert fake_file == os.path.join(scan_im_dir, "bar.img")
    # But this is none if we check it exists
    fake_file = scan_dir_manager.get_file_path_from_img_dir(
        scan_name, "bar.img", check_exists=True
    )
    assert fake_file is None

    # Create the dir
    scan_dir = scan_dir_manager.new_scan_dir("fake_scan")
    # This returns a ScanDirectory object
    assert isinstance(scan_dir, ScanDirectory)
    # The directory now exists as does the image directory
    assert scan_dir_manager.exists(scan_name)
    assert os.path.isdir(scan_path)
    assert os.path.isdir(scan_im_dir)

    # Test cleanup. Delete the created scan and scan directory
    scan_dir_manager.delete_scan(scan_name)
    assert not scan_dir_manager.exists(scan_name)
    assert not os.path.isdir(scan_path)


def test_bad_scan_names():
    """Check scan names with spaces, or worse BASH commands are not allowed."""
    _clear_scan_dir()
    scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR)
    scan_dir = scan_dir_manager.new_scan_dir("fake scan")
    assert scan_dir.name == "fake_scan_0001"
    scan_dir = scan_dir_manager.new_scan_dir("fake scan;rm -rf /;")
    assert scan_dir.name == "fake_scan_rm_-rf____0001"


def test_scan_sequence_and_listing():
    """Check created scans are added in order and listed correctly."""
    _clear_scan_dir()
    scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR)
    # Make 4 scans
    scan_dir_manager.new_scan_dir("fake_scan")
    scan_dir_manager.new_scan_dir("fake_scan")
    scan_dir_manager.new_scan_dir("fake_scan")
    scan_dir_manager.new_scan_dir("fake_scan")

    # Check they exist and are numbered sequentially
    all_scans = scan_dir_manager.all_scans
    assert_unique_of_length(all_scans, 4)
    assert "fake_scan_0001" in all_scans
    assert "fake_scan_0002" in all_scans
    assert "fake_scan_0003" in all_scans
    assert "fake_scan_0004" in all_scans

    # Check scan data can be read for all of them
    # (more detailed scan_info tests below)
    all_scan_info = scan_dir_manager.all_scans_info()
    for scan_info in all_scan_info:
        assert isinstance(scan_info, ScanInfo)
        assert scan_info.name.startswith("fake_scan_000")
        assert scan_info.number_of_images == 0


def test_scan_name_non_sequential():
    """Check new scan has the correct name if the directories are not sequential."""
    _clear_scan_dir()
    os.makedirs(os.path.join(BASE_SCAN_DIR, "fake_scan_0001"))
    os.makedirs(os.path.join(BASE_SCAN_DIR, "fake_scan_0002"))
    os.makedirs(os.path.join(BASE_SCAN_DIR, "fake_scan_0003"))
    os.makedirs(os.path.join(BASE_SCAN_DIR, "fake_scan_0005"))
    os.makedirs(os.path.join(BASE_SCAN_DIR, "fake_scan_0007"))
    os.makedirs(os.path.join(BASE_SCAN_DIR, "fake_scan_0011"))
    scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR)

    scan_dir = scan_dir_manager.new_scan_dir("fake_scan")
    assert scan_dir.name == "fake_scan_0012"

    all_scans = scan_dir_manager.all_scans
    assert_unique_of_length(all_scans, 7)
    assert "fake_scan_0001" in all_scans
    assert "fake_scan_0002" in all_scans
    assert "fake_scan_0003" in all_scans
    assert "fake_scan_0005" in all_scans
    assert "fake_scan_0007" in all_scans
    assert "fake_scan_0011" in all_scans
    assert "fake_scan_0012" in all_scans


def test_all_scan_names_taken():
    """If the next sequential scan name needs more than 4 digits check error is thrown."""
    _clear_scan_dir()
    os.makedirs(os.path.join(BASE_SCAN_DIR, "fake_scan_9999"))
    scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR)
    with pytest.raises(FileExistsError):
        scan_dir_manager.new_scan_dir("fake_scan")


def test_no_scan_names_given():
    """Check correct default scan name is used if empty string is given."""
    _clear_scan_dir()
    scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR)
    scan_dir = scan_dir_manager.new_scan_dir("")
    assert scan_dir.name == "scan_0001"


def test_scan_info():
    """Test the scan info is correct even using fake scan data."""
    _clear_scan_dir()
    scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR)
    scan_dir = scan_dir_manager.new_scan_dir("fake_scan")

    for _i in range(17):
        _add_fake_image(scan_dir)
    info = scan_dir.scan_info()
    now = time.time()

    assert info.name == "fake_scan_0001"
    # Created and modified in the last 5 seconds
    assert now - 5 < info.created < now
    assert now - 5 < info.modified < now
    assert info.number_of_images == 17
    assert not info.stitch_available
    assert info.dzi is None

    # Add a fake stitched images and check this is recognised as a stitch not
    # a scan image
    _add_fake_file(scan_dir, "fake_scan_0001_stitched.jpg", in_im_dir=True)
    info = scan_dir.scan_info()
    assert info.number_of_images == 17
    assert info.stitch_available


def test_get_final_stitch():
    """Check that the final stitch can be retrieved."""
    _clear_scan_dir()
    scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR)
    scan_dir = scan_dir_manager.new_scan_dir("fake_scan")

    # Create some scan images files
    for _i in range(17):
        _add_fake_image(scan_dir)

    # No scans, so None should be returned, from both manager and scan dir
    assert scan_dir_manager.get_final_stitch_path(scan_dir.name) is None
    assert scan_dir.get_final_stitch_name() is None

    fake_scan_name = "fake_scan_0001_stitched.jpg"
    fake_scan_path = scan_dir_manager.get_file_path_from_img_dir(
        scan_name=scan_dir.name,
        filename="fake_scan_0001_stitched.jpg",
        check_exists=False,
    )
    # Add a fake scan
    _add_fake_file(scan_dir, fake_scan_name, in_im_dir=True)
    # Manager returns full path
    assert scan_dir_manager.get_final_stitch_path(scan_dir.name) == fake_scan_path
    # ScanDirectory object returns just the filename
    assert scan_dir.get_final_stitch_name() == fake_scan_name


def test_get_scan_data_path():
    """Check that a scan data path behaves as expected."""
    _clear_scan_dir()
    scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR)
    scan_dir = scan_dir_manager.new_scan_dir("fake_scan")
    scan_name = scan_dir.name
    expected_path = os.path.join(scan_dir.images_dir, SCAN_DATA_FILENAME)
    # Asking the scan manager for the path will return None as there is no file.
    assert scan_dir_manager.get_scan_data_path(scan_name) is None
    # Asking the scan directly will return the location the file should be located
    # this is so it can be created.
    assert scan_dir.scan_data_path == expected_path

    # Remove the images directory.
    shutil.rmtree(scan_dir.images_dir)
    # When the images directory is deleted, internally `scan_dir.images_dir`
    # will return `None`. Check that `get_scan_data_path` copes with the `None`
    # return from `scan_dir.images_dir`.
    assert scan_dir_manager.get_scan_data_path(scan_name) is None


def test_get_scan_data_dict():
    """Check that the dictionary for the scan data is returned, or None if doesn't exist."""
    _clear_scan_dir()
    scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR)
    scan_dir = scan_dir_manager.new_scan_dir("fake_scan")
    scan_name = scan_dir.name
    # Doesn't yet exist
    assert scan_dir_manager.get_scan_data_dict(scan_name) is None

    fake_data = {"foo": 1, "bar": "foobar"}
    with open(scan_dir.scan_data_path, "w", encoding="utf-8") as json_file:
        json.dump(fake_data, json_file)

    # Should now be able to load this fake data from disk
    assert scan_dir_manager.get_scan_data_dict(scan_name) == fake_data

    # Check None is returned if the data cannot be read.
    with open(scan_dir.scan_data_path, "w", encoding="utf-8") as json_file:
        json_file.write("this is not json")
    assert scan_dir_manager.get_scan_data_dict(scan_name) is None


def test_empty_scan_info():
    """Test the scan info is correct even if the scan is empty."""
    _clear_scan_dir()
    scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR)
    scan_dir = scan_dir_manager.new_scan_dir("fake_scan")

    info = scan_dir.scan_info()
    now = time.time()

    assert info.name == "fake_scan_0001"
    # Created and modified in the last 5 seconds
    assert now - 5 < info.created < now
    assert now - 5 < info.modified < now
    assert info.number_of_images == 0
    assert not info.stitch_available
    assert info.dzi is None


def test_zipping_scan_data():
    """Test zipping the scan images with fake image data."""
    # Run twice, once calling the ScanDirectory directly,
    # Once calling the ScanDirectoryManager
    for caller in ["scan_dir", "manager"]:
        _clear_scan_dir()
        scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR)
        scan_dir = scan_dir_manager.new_scan_dir("fake_scan")

        # Create 21 fake scan images, a fake stitch, and a fake zip
        for _i in range(21):
            _add_fake_image(scan_dir)
        _add_fake_file(scan_dir, "fake_scan_0001_stitched.jpg", in_im_dir=True)
        _add_fake_file(scan_dir, "zipfile.zip")

        # This fake dzi should have loads of images. The DZI should not be zipped!
        _make_fake_dzi(scan_dir)

        # zip the directory without setting as the final version. It should only
        # zip the 21 scan images
        if caller == "scan_dir":
            zip_fname = scan_dir.zip_files()
        else:
            zip_fname = scan_dir_manager.zip_scan("fake_scan_0001")
        zip_files = get_files_in_zip(zip_fname)
        assert_unique_of_length(zip_files, 21)

        # Zip again with final version on and there should be 22 images
        if caller == "scan_dir":
            scan_dir.zip_files(final_version=True)
        else:
            scan_dir_manager.zip_scan("fake_scan_0001", final_version=True)
        zip_files = get_files_in_zip(zip_fname)
        assert_unique_of_length(get_files_in_zip(zip_fname), 22)
        # Check the zips are not in the zip
        for file in zip_files:
            assert not file.endswith(".zip")
            assert not file.endswith(".dzi")


def test_saving_and_loading_scan_data():
    """Test that scan data is saved and loaded as expected."""
    _clear_scan_dir()
    scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR)
    scan_dir = scan_dir_manager.new_scan_dir("fake_scan")
    scan_name = scan_dir.name

    # Should start without a scan data file.
    assert not os.path.isfile(scan_dir.scan_data_path)
    # Create
    scan_data_obj = _fake_scan_data()
    scan_dir.save_scan_data(scan_data_obj)
    # File should now exist
    assert os.path.isfile(scan_dir.scan_data_path)

    # Dump the scan json to a string an reload it
    # Note that more detailed checking of the dumping and loading of ScanData is in
    # tests/test_scan_data.py
    scan_data_dict = json.loads(scan_data_obj.model_dump_json())
    # What is loaded from file should be the same as from dumping and loading.
    assert scan_dir_manager.get_scan_data_dict(scan_name) == scan_data_dict


def test_saving_scan_data_error():
    """Test that saving scan data if there is no images directory raises FileNotFoundError."""
    _clear_scan_dir()
    scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR)
    scan_dir = scan_dir_manager.new_scan_dir("fake_scan")

    # Remove the images directory.
    shutil.rmtree(scan_dir.images_dir)
    # Should raise FileNotFoundError.
    with pytest.raises(FileNotFoundError):
        scan_dir.save_scan_data(_fake_scan_data())


def test_all_files():
    """Test all_files returns the path, and respects skipped directories."""
    _clear_scan_dir()
    scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR)
    scan_dir = scan_dir_manager.new_scan_dir("fake_scan")
    # This fake dzi should have loads of images. The DZI should not be zipped!
    _make_fake_dzi(scan_dir)
    all_files = scan_dir.all_files()

    # As standard for 8 layers there are 89 jpegs and 1 dzi file.
    assert_unique_of_length(all_files, 90)
    dzi_file = None
    # Check all files exist
    for file in all_files:
        assert os.path.exists(os.path.join(scan_dir.dir_path, file))
        if file.endswith(".dzi"):
            # There is only 1 dzi file, so this should be None
            assert dzi_file is None
            dzi_file = file
    # Once loop is complete there should be a dzi file
    assert dzi_file is not None
    # Get the dzi tile directory name
    dzi_dir = os.path.basename(dzi_file[:-4] + "_files")

    # Get all files skipping the dzi tile directory
    all_files = scan_dir.all_files(skip_dirs=dzi_dir)
    # There is now only one file
    assert len(all_files) == 1
    # It is the DZI file
    assert all_files[0] == dzi_file


def test_creating_scan_dir_for_missing_scan():
    """Check creating ScanDirectory object for a dir that doesn't exist fails."""
    _clear_scan_dir()
    with pytest.raises(FileNotFoundError):
        ScanDirectory("not_real_0001", BASE_SCAN_DIR)


def test_none_returned_for_missing_images_dir():
    """None should be returned for image dir path if images dir does not exist.

    By default images directories are created at the same time the scan directory
    is created. However, edge cases such as problems in deletions, microscopes
    with older scans on, etc can cause and empty scan directory, so it is handled
    explicitly.
    """
    _clear_scan_dir()
    os.makedirs(os.path.join(BASE_SCAN_DIR, "fake_scan_0001"))
    scan_dir = ScanDirectory("fake_scan_0001", BASE_SCAN_DIR)
    assert scan_dir.images_dir is None
    # Also check that get scan files returns and empty list
    assert scan_dir.get_scan_files() == []


def test_extracting_files():
    """Test the private _find_files method of ScanDirectories.

    Add files to directory and check expected returns.
    """
    _clear_scan_dir()
    os.makedirs(os.path.join(BASE_SCAN_DIR, "fake_scan_0001", "images"))
    scan_dir = ScanDirectory("fake_scan_0001", BASE_SCAN_DIR)

    # Starting all lists should be empty
    scan_files = scan_dir.get_scan_files()
    assert scan_dir._extract_scan_images(scan_files) == []
    assert scan_dir._extract_final_stitches(scan_files) == []
    assert scan_dir._extract_dzi_files(scan_files) == []

    # Add a number of images
    for _i in range(2321):
        _add_fake_image(scan_dir)

    scan_files = scan_dir.get_scan_files()
    assert_unique_of_length(scan_dir._extract_scan_images(scan_files), 2321)
    assert_unique_of_length(scan_dir._extract_final_stitches(scan_files), 0)
    assert_unique_of_length(scan_dir._extract_dzi_files(scan_files), 0)

    # Add and a stitched image
    _add_fake_file(scan_dir, "fake_scan_0001_stitched.jpg", in_im_dir=True)

    scan_files = scan_dir.get_scan_files()
    assert_unique_of_length(scan_dir._extract_scan_images(scan_files), 2321)
    assert_unique_of_length(scan_dir._extract_final_stitches(scan_files), 1)
    assert_unique_of_length(scan_dir._extract_dzi_files(scan_files), 0)

    _make_fake_dzi(scan_dir)

    # check totals are still correct after adding a dzi with lots of tiles.
    scan_files = scan_dir.get_scan_files()
    scan_images = scan_dir._extract_scan_images(scan_files)
    stitches = scan_dir._extract_final_stitches(scan_files)
    dzi_files = scan_dir._extract_dzi_files(scan_files)
    assert_unique_of_length(scan_images, 2321)
    assert_unique_of_length(stitches, 1)
    assert_unique_of_length(dzi_files, 1)

    # And check the names are as expected
    assert stitches[0] == "fake_scan_0001_stitched.jpg"
    assert dzi_files[0] == "fake_scan_0001.dzi"


DiskUsage = namedtuple("DiskUsage", ["total", "used", "free"])


@pytest.mark.parametrize("free_space", [500_000_001, 700_000_000, 5_000_000_000])
def test_disk_not_full(mocker, free_space):
    """Check no error thrown if disk has over 500MB of space."""
    total_space = 16_000_000_000
    # Mock the disk_usage
    mocker.patch(
        "shutil.disk_usage",
        return_value=DiskUsage(
            total=total_space, used=total_space - free_space, free=free_space
        ),
    )

    scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR)
    scan_dir_manager.check_free_disk_space()


@pytest.mark.parametrize("free_space", [100_000_000, 499_999_999])
def test_disk_full(mocker, free_space):
    """Check error thrown if disk has under 500MB of space."""
    total_space = 16_000_000_000
    # Mock the disk_usage
    mocker.patch(
        "shutil.disk_usage",
        return_value=DiskUsage(
            total=total_space, used=total_space - free_space, free=free_space
        ),
    )

    scan_dir_manager = ScanDirectoryManager(BASE_SCAN_DIR)
    with pytest.raises(NotEnoughFreeSpaceError):
        scan_dir_manager.check_free_disk_space()