test_flip_image.py

import unittest

import numpy as np

from csst_proto import flip_image, read_default_image
from csst_proto import flip_multiple_images_jl, flip_multiple_images_mp


class FlipImageTestCase(unittest.TestCase):
    def test_flip_image(self):
        """test flip image"""
        self.assertTrue(
            np.all(flip_image(read_default_image()) == np.array([[4, 3], [2, 1]]))
        )

        # the code fails for 1D array
        with self.assertRaises(AssertionError):
            flip_image(np.array([1, 2, 3, 4]))

    def test_flip_multiple_images_mp(self):
        """test flip multiple images with multiprocessing"""
        n_jobs = 10
        imgs = [read_default_image() for i_job in range(n_jobs)]
        flipped_imgs = flip_multiple_images_mp(imgs, n_jobs)
        for i_job in range(n_jobs):
            self.assertTrue(np.all(flipped_imgs[i_job] == np.array([[4, 3], [2, 1]])))

    def test_flip_multiple_images_jl(self):
        """test flip multiple images with joblib"""
        n_jobs = 10
        imgs = [read_default_image() for i_job in range(n_jobs)]
        flipped_imgs = flip_multiple_images_jl(imgs, n_jobs)
        for i_job in range(n_jobs):
            self.assertTrue(np.all(flipped_imgs[i_job] == np.array([[4, 3], [2, 1]])))