import io
import os
import time
import types
import unittest
import warnings
from contextlib import redirect_stderr, redirect_stdout
from itertools import product, repeat
from multiprocessing import Barrier, Value
from threading import current_thread, main_thread, Thread
from unittest.mock import Mock, patch

import numpy as np
from tqdm import tqdm

from mpire import cpu_count, WorkerPool
from mpire.async_result import AsyncResult
from mpire.context import FORK_AVAILABLE, RUNNING_WINDOWS

# Skip start methods that use fork if it's not available
if not FORK_AVAILABLE:
    TEST_START_METHODS = ['spawn', 'threading']
else:
    TEST_START_METHODS = ['fork', 'forkserver', 'spawn', 'threading']


def square(idx, x):
    return idx, x * x


def extremely_large_output(idx, _):
    return idx, os.urandom(1024 * 1024)


def square_numpy(x):
    return x * x


def subtract(x, y):
    return x - y


class MapTest(unittest.TestCase):

    def setUp(self):
        # Create some test data. Note that the regular map reads the inputs as a list of single tuples (one argument),
        # whereas parallel.map sees it as a list of argument lists. Therefore we give the regular map a lambda function
        # which mimics the parallel.map behavior.
        self.test_data = list(enumerate([1, 2, 3, 5, 6, 9, 37, 42, 1337, 0, 3, 5, 0]))
        self.test_desired_output = list(map(lambda _args: square(*_args), self.test_data))
        self.test_data_len = len(self.test_data)

        # Numpy test data
        self.test_data_numpy = np.random.rand(100, 2)
        self.test_desired_output_numpy = square_numpy(self.test_data_numpy)
        self.test_data_len_numpy = len(self.test_data_numpy)

    def test_all_maps(self):
        """
        Tests the map related functions
        """
        def get_generator(iterable):
            yield from iterable

        # Test results for different parameter settings
        print()
        for n_jobs, n_tasks_max_active, worker_lifespan, chunk_size, n_splits in tqdm([
            (None, None, None, None, None),
            (1, None, None, None, None),
            (2, None, None, None, None),
            (2, 2, None, None, None),
            (2, None, 2, None, None),
            (2, None, None, 3, None),
            (2, None, None, None, 3),
            (2, None, None, 3, 3),
            (2, None, 1, 3, None)
        ]):
            with WorkerPool(n_jobs=n_jobs) as pool:

                for map_func, sort, result_type in ((pool.map, False, list), (pool.map_unordered, True, list),
                                                    (pool.imap, False, types.GeneratorType),
                                                    (pool.imap_unordered, True, types.GeneratorType)):

                    with self.subTest(map_func=map_func, input='list', n_jobs=n_jobs,
                                      n_tasks_max_active=n_tasks_max_active, worker_lifespan=worker_lifespan,
                                      chunk_size=chunk_size, n_splits=n_splits):

                        # Test if parallel map results in the same as ordinary map function. Should work both for
                        # generators and iterators. Also check if an empty list and extremely large output (exceeding
                        # os.pipe limits) works as desired.
                        results_list = map_func(square, self.test_data, max_tasks_active=n_tasks_max_active,
                                                worker_lifespan=worker_lifespan)
                        self.assertIsInstance(results_list, result_type)
                        self.assertEqual(self.test_desired_output,
                                         sorted(results_list, key=lambda tup: tup[0]) if sort else list(results_list))

                    with self.subTest(map_func=map_func, input='generator', n_jobs=n_jobs,
                                      n_tasks_max_active=n_tasks_max_active, worker_lifespan=worker_lifespan,
                                      chunk_size=chunk_size, n_splits=n_splits):

                        results_list = map_func(square, get_generator(self.test_data), iterable_len=self.test_data_len,
                                                max_tasks_active=n_tasks_max_active, worker_lifespan=worker_lifespan)
                        self.assertIsInstance(results_list, result_type)
                        self.assertEqual(self.test_desired_output,
                                         sorted(results_list, key=lambda tup: tup[0]) if sort else list(results_list))

                    with self.subTest(map_func=map_func, input='empty list', n_jobs=n_jobs,
                                      n_tasks_max_active=n_tasks_max_active, worker_lifespan=worker_lifespan,
                                      chunk_size=chunk_size, n_splits=n_splits):

                        results_list = map_func(square, [], max_tasks_active=n_tasks_max_active,
                                                worker_lifespan=worker_lifespan)
                        self.assertIsInstance(results_list, result_type)
                        self.assertEqual([], list(results_list))

                    # When the os pipe capacity is exceeded, a worker restart based on worker lifespan would hang if we
                    # not fetch all the results from a worker. We only verify the amount of data returned here.
                    with self.subTest(map_func=map_func, output='data exceeding pipe limits', n_jobs=n_jobs,
                                      n_tasks_max_active=n_tasks_max_active, worker_lifespan=worker_lifespan,
                                      chunk_size=chunk_size, n_splits=n_splits):
                        results_list = map_func(extremely_large_output, self.test_data,
                                                max_tasks_active=n_tasks_max_active, worker_lifespan=worker_lifespan)
                        self.assertIsInstance(results_list, result_type)
                        self.assertEqual(len(self.test_desired_output), len(list(results_list)))

    def test_numpy_input(self):
        """
        Test map with numpy input
        """
        print()
        for n_jobs, n_tasks_max_active, worker_lifespan, chunk_size, n_splits in tqdm([
            (None, None, None, None, None),
            (1, None, None, None, None),
            (2, None, None, None, None),
            (2, 2, None, None, None),
            (2, None, 2, None, None),
            (2, None, None, 3, None),
            (2, None, None, None, 3),
            (2, None, None, 3, 3),
            (2, None, 1, 3, None)
        ]):
            with WorkerPool(n_jobs=n_jobs) as pool:

                # Test numpy input. map should concatenate chunks of numpy output to a single output array if we
                # instruct it to
                with self.subTest(concatenate_numpy_output=True, map_function='map', n_jobs=n_jobs,
                                  n_tasks_max_active=n_tasks_max_active, worker_lifespan=worker_lifespan,
                                  chunk_size=chunk_size, n_splits=n_splits):
                    results = pool.map(square_numpy, self.test_data_numpy, max_tasks_active=n_tasks_max_active,
                                       worker_lifespan=worker_lifespan, concatenate_numpy_output=True)
                    self.assertIsInstance(results, np.ndarray)
                    np.testing.assert_array_equal(results, self.test_desired_output_numpy)

                # If we disable it we should get back chunks of the original array
                with self.subTest(concatenate_numpy_output=False, map_function='map', n_jobs=n_jobs,
                                  n_tasks_max_active=n_tasks_max_active, worker_lifespan=worker_lifespan,
                                  chunk_size=chunk_size, n_splits=n_splits):
                    results = pool.map(square_numpy, self.test_data_numpy, max_tasks_active=n_tasks_max_active,
                                       worker_lifespan=worker_lifespan, concatenate_numpy_output=False)
                    self.assertIsInstance(results, list)
                    np.testing.assert_array_equal(np.concatenate(results), self.test_desired_output_numpy)

                # Numpy concatenation doesn't exist for the other functions
                with self.subTest(map_function='imap', n_jobs=n_jobs, n_tasks_max_active=n_tasks_max_active,
                                  worker_lifespan=worker_lifespan, chunk_size=chunk_size, n_splits=n_splits):
                    results = pool.imap(square_numpy, self.test_data_numpy, max_tasks_active=n_tasks_max_active,
                                        worker_lifespan=worker_lifespan)
                    self.assertIsInstance(results, types.GeneratorType)
                    np.testing.assert_array_equal(np.concatenate(list(results)), self.test_desired_output_numpy)

                # map_unordered and imap_unordered cannot be checked for correctness as we don't know the order of the
                # returned results, except when n_jobs=1. In the other cases we could, however, check if all the values
                # (numpy rows) that are returned are present (albeit being in a different order)
                for map_func, result_type in ((pool.map_unordered, list), (pool.imap_unordered, types.GeneratorType)):

                    with self.subTest(map_function=map_func, n_jobs=n_jobs, n_tasks_max_active=n_tasks_max_active,
                                      worker_lifespan=worker_lifespan, chunk_size=chunk_size, n_splits=n_splits):

                        results = map_func(square_numpy, self.test_data_numpy, max_tasks_active=n_tasks_max_active,
                                           worker_lifespan=worker_lifespan)
                        self.assertIsInstance(results, result_type)
                        concattenated_results = np.concatenate(list(results))
                        if n_jobs == 1:
                            np.testing.assert_array_equal(concattenated_results, self.test_desired_output_numpy)
                        else:
                            # We sort the expected and actual results using lexsort, which sorts using a sequence of
                            # keys. We transpose the array to sort on columns instead of rows.
                            np.testing.assert_array_equal(
                                concattenated_results[np.lexsort(concattenated_results.T)],
                                self.test_desired_output_numpy[np.lexsort(self.test_desired_output_numpy.T)]
                            )

    def test_dictionary_input(self):
        """
        Test map with dictionary input
        """
        with WorkerPool(n_jobs=1) as pool:

            # Should work
            with self.subTest('correct input'):
                results_list = pool.map(subtract, [{'x': 5, 'y': 2}, {'y': 5, 'x': 2}])
                self.assertEqual(results_list, [3, -3])

            # Should throw
            with self.subTest("missing 'y', unknown parameter 'z'"), self.assertRaises(TypeError):
                pool.map(subtract, [{'x': 5, 'z': 2}])

            # Should throw
            with self.subTest("unknown parameter 'z'"), self.assertRaises(TypeError):
                pool.map(subtract, [{'x': 5, 'y': 2, 'z': 2}])

    def test_start_methods(self):
        """
        Test different start methods. All should work just fine
        """
        print()
        for start_method in tqdm(TEST_START_METHODS):
            with self.subTest(start_method=start_method, map='map'), WorkerPool(2, start_method=start_method) as pool:
                results_list = pool.map(square, self.test_data)
                self.assertIsInstance(results_list, list)
                self.assertEqual(self.test_desired_output, results_list)

            with self.subTest(start_method=start_method, map='map_unordered'), \
                    WorkerPool(2, start_method=start_method) as pool:
                results_list = pool.map_unordered(square, self.test_data)
                self.assertIsInstance(results_list, list)
                self.assertEqual(self.test_desired_output, sorted(results_list, key=lambda tup: tup[0]))

            with self.subTest(start_method=start_method, map='imap'), WorkerPool(2, start_method=start_method) as pool:
                results_list = pool.imap(square, self.test_data)
                self.assertIsInstance(results_list, types.GeneratorType)
                self.assertListEqual(list(results_list), self.test_desired_output)

            with self.subTest(start_method=start_method, map='imap_unordered'), \
                    WorkerPool(2, start_method=start_method) as pool:
                results_list = pool.imap_unordered(square, self.test_data)
                self.assertIsInstance(results_list, types.GeneratorType)
                self.assertEqual(self.test_desired_output, sorted(results_list, key=lambda tup: tup[0]))

    def test_mixing_map_calls(self):
        """
        When using the same pool, mixing map calls should raise
        """
        with WorkerPool(2) as pool:
            imap_results = pool.imap(square, self.test_data)
            next(imap_results)  # Actually start the pool
            with self.assertRaises(RuntimeError):
                pool.map(square, self.test_data)

        with WorkerPool(2) as pool:
            imap_results = pool.imap_unordered(square, self.test_data)
            next(imap_results)  # Actually start the pool
            with self.assertRaises(RuntimeError):
                next(pool.imap(square, self.test_data))

    def test_terminate(self):
        """
        When a lazy map call is running and the pool is terminated, exhausting the results should raise
        """
        with self.subTest("calling terminate() explicitly"), WorkerPool(1) as pool:
            imap_results = pool.imap(square, self.test_data)
            next(imap_results)  # Actually start the pool
            pool.terminate()
            with self.assertRaises(RuntimeError):
                list(imap_results)

        with self.subTest("calling terminate() implicitly"):
            with WorkerPool(1) as pool:
                imap_results = pool.imap(square, self.test_data)
                next(imap_results)  # Actually start the pool
            with self.assertRaises(RuntimeError):
                list(imap_results)

        # Before, this could cause a deadlock once all tests were done
        print()
        with self.subTest("calling terminate() implicitly, with progress bar"):
            with WorkerPool(1) as pool:
                imap_results = pool.imap(square, self.test_data, progress_bar=True)
                next(imap_results)  # Actually start the pool
            with self.assertRaises(RuntimeError):
                list(imap_results)


class PoolInThreadTest(unittest.TestCase):

    def setUp(self):
        self.test_data = [1, 2, 3, 5, 6, 9, 37, 42, 1337, 0, 3, 5, 0]
        self.test_desired_output = [self._square(x) for x in self.test_data]

    def test_start_methods(self):
        """
        Test that a WorkerPool can be started inside a thread, which isn't the main thread. Test for different start
        methods. All should work just fine
        """
        for start_method in TEST_START_METHODS:
            with self.subTest(start_method=start_method):
                t = Thread(target=self._map_thread, args=(start_method,))
                t.start()
                t.join()

    def _map_thread(self, start_method):
        """
        This function is called from within a thread
        """
        self.assertNotEqual(current_thread(), main_thread())
        with WorkerPool(2, start_method=start_method) as pool:
            results_list = pool.map(self._square, self.test_data)
            self.assertIsInstance(results_list, list)
            self.assertListEqual(self.test_desired_output, results_list)

    @staticmethod
    def _square(x):
        return x * x


class ApplyTest(unittest.TestCase):

    def test_apply_async_call(self):
        """
        Test that apply simply calls apply_async
        """
        with WorkerPool(1) as pool, patch.object(pool, 'apply_async') as mock_apply_async:
            pool.apply(subtract, (1,), {'y': 2}, self._callback, self._error_callback,
                       self._init, self._exit, 0.1, 0.2, 0.3)
            mock_apply_async.assert_called_once_with(subtract, (1,), {'y': 2}, self._callback, self._error_callback,
                                                     self._init, self._exit, 0.1, 0.2, 0.3)

    def test_result(self):
        """
        Test that apply returns the correct result
        """
        with WorkerPool(1) as pool:
            results = [pool.apply(self._square, (i,)) for i in range(10)]
            self.assertEqual(results, [self._square(i) for i in range(10)])

    @staticmethod
    def _callback(_):
        return 0

    @staticmethod
    def _error_callback(_):
        return 1

    @staticmethod
    def _init():
        return

    @staticmethod
    def _exit():
        return 2

    @staticmethod
    def _square(x):
        return x * x


class ApplyAsyncTest(unittest.TestCase):

    def test_result(self):
        """
        Test that apply_async returns the correct result. Calling get multiple times should also work
        """
        with WorkerPool(2) as pool:
            results = [pool.apply_async(self._square, (i,)) for i in range(10)]
            [self.assertIsInstance(result, AsyncResult) for result in results]
            self.assertListEqual([result.get() for result in results], [self._square(i) for i in range(10)])
            self.assertListEqual([result.get() for result in results], [self._square(i) for i in range(10)])

    def test_args_kwargs(self):
        """
        Test that apply_async works with args and kwargs
        """
        with WorkerPool(2) as pool:
            results = [pool.apply_async(subtract, (i * i,), {'y': i}) for i in range(10)]
            self.assertListEqual([result.get() for result in results], [subtract(i * i, i) for i in range(10)])

            results = [pool.apply_async(subtract, (), {'x': i * i, 'y': i}) for i in range(10)]
            self.assertListEqual([result.get() for result in results], [subtract(i * i, i) for i in range(10)])

    def test_callback(self):
        """
        Test that apply_async calls the callback function on success
        """
        callback = Mock()
        with WorkerPool(1) as pool:
            pool.apply_async(self._square, (42,), callback=callback).get()
            callback.assert_called_once_with(42 * 42)

    def test_callback_error(self):
        """
        Test that apply_async calls the error callback function on error
        """
        callback = Mock()
        with WorkerPool(1) as pool:
            value_error = ValueError('test')
            with self.assertRaises(ValueError):
                pool.apply_async(self._raise_exception, (value_error,), error_callback=callback).get()
            self.assertIsInstance(callback.call_args[0][0], ValueError)

    def test_second_apply_raises(self):
        """
        When a second apply task raises an exception, the first task should still be able to complete. I.e., the second
        worker shouldn't cause the entire pool to shutdown
        """
        with self.subTest("exception is raised"), WorkerPool(2) as pool:
            event = pool.ctx.Event()
            pool.set_shared_objects(event)
            first_result = pool.apply_async(self._wait_and_return, (42,))
            with self.assertRaises(ValueError):
                pool.apply_async(self._raise_exception_2).get()
            self.assertFalse(first_result.ready())
            event.set()
            self.assertEqual(first_result.get(), 42)

        with self.subTest("timeout is raised"), WorkerPool(2) as pool:
            event = pool.ctx.Event()
            pool.set_shared_objects(event)
            first_result = pool.apply_async(self._wait_and_return, (42,))
            with self.assertRaises(TimeoutError):
                pool.apply_async(self._wait_and_return, (1337,), task_timeout=0.01).get()
            self.assertFalse(first_result.ready())
            event.set()
            self.assertEqual(first_result.get(), 42)

    @staticmethod
    def _square(x):
        return x * x

    @staticmethod
    def _raise_exception(exception):
        raise exception

    @staticmethod
    def _raise_exception_2(_):
        raise ValueError

    @staticmethod
    def _wait_and_return(e, x):
        e.wait()
        return x


class WorkerIDTest(unittest.TestCase):

    def test_by_config_function(self):
        """
        Test setting passing on the worker ID using the pass_on_worker_id function
        """
        for n_jobs, pass_worker_id in product([1, 3], [True, False]):

            with self.subTest(n_jobs=n_jobs, pass_worker_id=pass_worker_id, config_type='function'), \
                 WorkerPool(n_jobs=n_jobs) as pool:

                pool.pass_on_worker_id(pass_worker_id)

                # Tests should fail when number of arguments in function is incorrect, worker ID is not within range,
                # or when the shared objects are not equal to the given arguments
                f = self._f1 if pass_worker_id else self._f2
                self.assertListEqual(pool.map(f, ((n_jobs,) for _ in range(10)), iterable_len=10), [True] * 10)

    def test_by_constructor(self):
        """
        Test setting passing on the worker ID in the constructor
        """
        for n_jobs, pass_worker_id in product([1, 3], [True, False]):

            with self.subTest(n_jobs=n_jobs, pass_worker_id=pass_worker_id, config_type='constructor'), \
                 WorkerPool(n_jobs=n_jobs, pass_worker_id=pass_worker_id) as pool:

                # Tests should fail when number of arguments in function is incorrect, worker ID is not within range,
                # or when the shared objects are not equal to the given arguments
                f = self._f1 if pass_worker_id else self._f2
                self.assertListEqual(pool.map(f, ((n_jobs,) for _ in range(10)), iterable_len=10), [True] * 10)

    def test_start_methods(self):
        """
        Test for different start methods
        """
        for start_method in TEST_START_METHODS:
            with self.subTest(start_method=start_method), \
                    WorkerPool(n_jobs=2, pass_worker_id=True, start_method=start_method) as pool:
                self.assertListEqual(pool.map(self._f1, ((2,) for _ in range(10)), iterable_len=10), [True] * 10)

    @staticmethod
    def _f1(_wid, _n_jobs):
        """
        Function with worker ID
        """
        tests_succeed = True
        tests_succeed &= isinstance(_wid, int)
        tests_succeed &= _wid >= 0
        tests_succeed &= _wid <= _n_jobs
        return tests_succeed

    @staticmethod
    def _f2(_n_jobs):
        """
        Function without worker ID (simply tests if WorkerPool correctly handles pass_worker_id=False)
        """
        return True


class SharedObjectsTest(unittest.TestCase):

    def test_by_config_function(self):
        """
        Tests passing shared objects using the set_shared_objects function
        """
        for n_jobs, shared_objects in product([1, 3], [None, (37, 42), ({'1', '2', '3'})]):

            with self.subTest(n_jobs=n_jobs, shared_objects=shared_objects, config_type='function'), \
                 WorkerPool(n_jobs=n_jobs) as pool:

                # Configure pool
                pool.set_shared_objects(shared_objects)

                # Tests should fail when number of arguments in function is incorrect, worker ID is not within range,
                # or when the shared objects are not equal to the given arguments
                f = self._f1 if shared_objects else self._f2
                self.assertListEqual(pool.map(f, ((shared_objects,) for _ in range(10)), iterable_len=10), [True] * 10)

    def test_by_constructor(self):
        """
        Tests passing shared objects in the constructor
        """
        for n_jobs, shared_objects in product([1, 3], [None, (37, 42), ({'1', '2', '3'})]):

            # Pass on arguments using the constructor instead
            with self.subTest(n_jobs=n_jobs, shared_objects=shared_objects, config_type='constructor'), \
                 WorkerPool(n_jobs=n_jobs, shared_objects=shared_objects) as pool:

                # Tests should fail when number of arguments in function is incorrect, worker ID is not within range,
                # or when the shared objects are not equal to the given arguments
                f = self._f1 if shared_objects else self._f2
                self.assertListEqual(pool.map(f, ((shared_objects,) for _ in range(10)), iterable_len=10), [True] * 10)

    def test_start_methods(self):
        """
        Tests for different start methods
        """
        for start_method in TEST_START_METHODS:
            with self.subTest(start_method=start_method), \
                    WorkerPool(n_jobs=2, shared_objects=({'1', '2', '3'}), start_method=start_method) as pool:
                self.assertListEqual(pool.map(self._f1, (({'1', '2', '3'},) for _ in range(10)), iterable_len=10),
                                     [True] * 10)

    @staticmethod
    def _f1(_sobjects, _args):
        """
        Function with shared objects
        """
        return _sobjects == _args

    @staticmethod
    def _f2(_args):
        """
        Function without shared objects (simply tests if WorkerPool correctly handles shared_objects=None)
        """
        return True


class WorkerStateTest(unittest.TestCase):

    def test_by_config_function(self):
        """
        Tests setting worker state using the set_use_worker_state function
        """
        for n_jobs, use_worker_state, n_tasks in product([1, 3], [False, True], [0, 1, 150]):

            with self.subTest(n_jobs=n_jobs, use_worker_state=use_worker_state, n_tasks=n_tasks),\
                 WorkerPool(n_jobs=n_jobs, pass_worker_id=True) as pool:

                pool.set_use_worker_state(use_worker_state)

                # When use_worker_state is set, the final (worker_id, n_args) of each worker should add up to the
                # number of given tasks
                f = self._f1 if use_worker_state else self._f2
                results = pool.map(f, range(n_tasks), chunk_size=2)
                if use_worker_state:
                    n_processed_per_worker = [0] * n_jobs
                    for wid, n_processed, tests_succeed in results:
                        n_processed_per_worker[wid] = n_processed
                        self.assertTrue(tests_succeed)
                    self.assertEqual(sum(n_processed_per_worker), n_tasks)

    def test_by_constructor(self):
        """
        Tests setting worker state in the constructor
        """
        for n_jobs, use_worker_state, n_tasks in product([1, 3], [False, True], [0, 1, 150]):

            with self.subTest(n_jobs=n_jobs, use_worker_state=use_worker_state, n_tasks=n_tasks), \
                 WorkerPool(n_jobs=n_jobs, pass_worker_id=True, use_worker_state=use_worker_state) as pool:

                # When use_worker_state is set, the final (worker_id, n_args) of each worker should add up to the
                # number of given tasks
                f = self._f1 if use_worker_state else self._f2
                results = pool.map(f, range(n_tasks), chunk_size=2)
                if use_worker_state:
                    n_processed_per_worker = [0] * n_jobs
                    for wid, n_processed, tests_succeed in results:
                        n_processed_per_worker[wid] = n_processed
                        self.assertTrue(tests_succeed)
                    self.assertEqual(sum(n_processed_per_worker), n_tasks)

    def test_start_methods(self):
        """
        Test for different start methods
        """
        for start_method in TEST_START_METHODS:
            with self.subTest(start_method=start_method), \
                    WorkerPool(n_jobs=2, pass_worker_id=True, use_worker_state=True, start_method=start_method) as pool:
                results = pool.map(self._f1, range(10), chunk_size=2)
                n_processed_per_worker = [0, 0, 0]
                for wid, n_processed, tests_succeed in results:
                    n_processed_per_worker[wid] = n_processed
                    self.assertTrue(tests_succeed)
                self.assertEqual(sum(n_processed_per_worker), 10)

    @staticmethod
    def _f1(_wid, _wstate, _arg):
        """
        Function with worker ID and worker state
        """
        tests_succeed = True
        tests_succeed &= isinstance(_wstate, dict)

        # Worker id should always be the same
        _wstate.setdefault('worker_id', set()).add(_wid)
        tests_succeed &= _wstate['worker_id'] == {_wid}

        # Should contain previous args
        _wstate.setdefault('args', []).append(_arg)
        return _wid, len(_wstate['args']), tests_succeed

    @staticmethod
    def _f2(_wid, _):
        """
        Function with worker ID (simply tests if WorkerPool correctly handles use_worker_state=False)
        """
        pass


class InitFuncTest(unittest.TestCase):

    def setUp(self) -> None:
        self.test_data = range(10)
        self.test_desired_output = [42, 43, 44, 45, 46, 47, 48, 49, 50, 51]

    def test_no_init_func(self):
        """
        If the init func is not provided, then `worker_state['test']` should fail
        """
        with self.assertRaises(KeyError), WorkerPool(n_jobs=4, shared_objects=(None,), use_worker_state=True) as pool:
            pool.map(self._f, range(10), worker_init=None)

    def test_init_func(self):
        """
        Test if init func is called. If it is, then `worker_state['test']` should be available. Due to the barrier we
        know for sure that the init func should be called as many times as there are workers
        """
        for n_jobs in [1, 3]:
            shared_objects = Barrier(n_jobs), Value('i', 0)
            with self.subTest(n_jobs=n_jobs), WorkerPool(n_jobs=n_jobs, shared_objects=shared_objects,
                                                         use_worker_state=True) as pool:
                results = pool.map(self._f, self.test_data, worker_init=self._init, chunk_size=1)
                self.assertListEqual(results, self.test_desired_output)
                self.assertEqual(shared_objects[1].value, n_jobs)

    def test_worker_lifespan(self):
        """
        When workers have a limited lifespan they are spawned multiple times. Each time a worker starts it should call
        the init function. Due to the chunk size we know for sure that the init func should be called at least once for
        each task. However, when all tasks have been processed the workers are terminated and we don't know exactly how
        many workers restarted. We only know for sure that the init func should be called between 10 and 10 + n_jobs
        times
        """
        for n_jobs in [1, 3]:
            shared_objects = Barrier(n_jobs), Value('i', 0)
            with self.subTest(n_jobs=n_jobs), WorkerPool(n_jobs=n_jobs, shared_objects=shared_objects,
                                                         use_worker_state=True) as pool:
                results = pool.map(self._f, self.test_data, worker_init=self._init, chunk_size=1, worker_lifespan=1)
                self.assertListEqual(results, self.test_desired_output)
                self.assertGreaterEqual(shared_objects[1].value, 10)
                self.assertLessEqual(shared_objects[1].value, 10 + n_jobs)

    def test_error(self):
        """
        When an exception occurs in the init function it should properly shut down
        """
        with self.subTest("map"), self.assertRaises(ValueError), \
                WorkerPool(n_jobs=4, shared_objects=(None,), use_worker_state=True) as pool:
            pool.map(self._f, self.test_data, worker_init=self._init_error)

        with self.subTest("apply"), self.assertRaises(ValueError), \
                WorkerPool(n_jobs=2, shared_objects=(None,), use_worker_state=True) as pool:
            pool.apply(self._f, args=(0,), worker_init=self._init_error)

    def test_start_methods(self):
        """
        Test for different start methods
        """
        for start_method in TEST_START_METHODS:
            with self.subTest(start_method=start_method), \
                    WorkerPool(n_jobs=2, use_worker_state=True, start_method=start_method) as pool:
                shared_objects = pool.ctx.Barrier(2), pool.ctx.Value('i', 0)
                pool.set_shared_objects(shared_objects)
                results = pool.map(self._f, self.test_data, worker_init=self._init, chunk_size=1)
                self.assertListEqual(results, self.test_desired_output)
                self.assertEqual(shared_objects[1].value, 2)

    @staticmethod
    def _init(shared_objects, worker_state):
        barrier, call_count = shared_objects

        # Only wait for the other workers the first time around (it will hang when worker_lifespan=1, otherwise)
        if call_count.value == 0:
            barrier.wait()

        with call_count.get_lock():
            call_count.value += 1
        worker_state['test'] = 42

    @staticmethod
    def _init_error(*_):
        raise ValueError(":(")

    @staticmethod
    def _f(_, worker_state, x):
        return worker_state['test'] + x


class ExitFuncTest(unittest.TestCase):

    def setUp(self) -> None:
        self.test_data = range(10)
        self.test_desired_output = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

    def test_no_exit_func(self):
        """
        If the exit func is not provided, then exit results shouldn't be available
        """
        shared_objects = Barrier(4), Value('i', 0)
        with WorkerPool(n_jobs=4, shared_objects=shared_objects, use_worker_state=True) as pool:
            results = pool.map(self._f1, range(10), worker_init=self._init, worker_exit=None)
            self.assertListEqual(results, self.test_desired_output)
            self.assertListEqual(pool.get_exit_results(), [])

    def test_exit_func(self):
        """
        Test if exit func is called. If it is, then exit results should be available. It should have as many elements
        as the number of jobs and should have the right content.
        """
        for n_jobs in [1, 3]:
            shared_objects = Barrier(n_jobs), Value('i', 0)
            with self.subTest(n_jobs=n_jobs), WorkerPool(n_jobs=n_jobs, shared_objects=shared_objects,
                                                         use_worker_state=True) as pool:
                results = pool.map(self._f1, self.test_data, worker_init=self._init, worker_exit=self._exit)
                self.assertListEqual(results, self.test_desired_output)
                self.assertEqual(shared_objects[1].value, n_jobs)
                self.assertEqual(len(pool.get_exit_results()), n_jobs)
                self.assertEqual(sum(pool.get_exit_results()), sum(range(10)))

    def test_worker_lifespan(self):
        """
        When workers have a limited lifespan they are spawned multiple times. Each time a worker exits it should call
        the exit function. Due to the chunk size we know for sure that the exit func should be called at least once for
        each task. However, when all tasks have been processed the workers are terminated and we don't know exactly how
        many workers restarted. We only know for sure that the exit func should be called between 10 and 10 + n_jobs
        times
        """
        for n_jobs in [1, 3]:
            shared_objects = Barrier(n_jobs), Value('i', 0)
            with self.subTest(n_jobs=n_jobs), WorkerPool(n_jobs=n_jobs, shared_objects=shared_objects,
                                                         use_worker_state=True) as pool:
                results = pool.map(self._f1, self.test_data, worker_init=self._init, worker_exit=self._exit,
                                   chunk_size=1, worker_lifespan=1)
                self.assertListEqual(results, self.test_desired_output)
                self.assertGreaterEqual(shared_objects[1].value, 10)
                self.assertLessEqual(shared_objects[1].value, 10 + n_jobs)
                self.assertEqual(len(pool.get_exit_results()), shared_objects[1].value)
                self.assertEqual(sum(pool.get_exit_results()), sum(range(10)))

    def test_exit_func_big_payload(self):
        """
        Multiprocessing Pipes have a maximum buffer size (depending on the system it can be anywhere between 16-1024kb).
        Results from the pipe need to be received from the other end, before the workers are joined. Otherwise the
        process can hang indefinitely. Because exit results are fetched in a different way as regular results, we test
        that here. We send a payload of 10_000kb.
        """
        for n_jobs, worker_lifespan in product([1, 3], [None, 2]):
            with self.subTest(n_jobs=n_jobs, worker_lifespan=worker_lifespan), WorkerPool(n_jobs=n_jobs) as pool:
                results = pool.map(self._f2, self.test_data, worker_exit=self._exit_big_payloud, chunk_size=1,
                                   worker_lifespan=worker_lifespan)
                self.assertListEqual(results, self.test_desired_output)
                self.assertTrue(bool(pool.get_exit_results()))
                for exit_result in pool.get_exit_results():
                    self.assertEqual(len(exit_result), 10_000 * 1024)

    def test_error(self):
        """
        When an exception occurs in the exit function it should properly shut down
        """
        for worker_lifespan in [None, 2]:
            with self.subTest("map", worker_lifespan=worker_lifespan), self.assertRaises(ValueError), \
                    WorkerPool(n_jobs=4) as pool:
                pool.map(self._f2, range(10), worker_lifespan=worker_lifespan, worker_exit=self._exit_error)

        with self.subTest("apply"), self.assertRaises(ValueError), WorkerPool(n_jobs=2) as pool:
            pool.apply(self._f2, args=(0,), worker_exit=self._exit_error)
            pool.stop_and_join()

    def test_start_methods(self):
        """
        Test for different start methods
        """
        for start_method in TEST_START_METHODS:
            with self.subTest(start_method=start_method), \
                    WorkerPool(n_jobs=2, use_worker_state=True, start_method=start_method) as pool:
                shared_objects = pool.ctx.Barrier(2), pool.ctx.Value('i', 0)
                pool.set_shared_objects(shared_objects)
                results = pool.map(self._f1, self.test_data, worker_init=self._init, worker_exit=self._exit)
                self.assertListEqual(results, self.test_desired_output)
                self.assertEqual(shared_objects[1].value, 2)
                self.assertEqual(len(pool.get_exit_results()), 2)
                self.assertEqual(sum(pool.get_exit_results()), sum(range(10)))

    @staticmethod
    def _init(shared_objects, worker_state):
        barrier, call_count = shared_objects

        # Only wait for the other workers the first time around (it will hang when worker_lifespan=1, otherwise)
        if call_count.value == 0:
            barrier.wait()

        worker_state['count'] = 0

    @staticmethod
    def _f1(_, worker_state, x):
        worker_state['count'] += x
        return x

    @staticmethod
    def _f2(x):
        return x

    @staticmethod
    def _exit(shared_objects, worker_state):
        _, call_count = shared_objects
        with call_count.get_lock():
            call_count.value += 1
        return worker_state['count']

    @staticmethod
    def _exit_big_payloud():
        return np.random.bytes(10_000 * 1024)

    @staticmethod
    def _exit_error():
        raise ValueError(":'(")


class DaemonTest(unittest.TestCase):

    # This also tests nested WorkerPools. We only test spawn here as creating processes is not thread-safe

    def setUp(self):
        # Create some test data. Note that the regular map reads the inputs as a list of single tuples (one argument),
        # whereas parallel.map sees it as a list of argument lists. Therefore we give the regular map a lambda function
        # which mimics the parallel.map behavior.
        self.test_data = list(enumerate([1, 2, 3, 5, 6, 9, 37, 42, 1337, 0, 3, 5, 0]))
        self.test_desired_output = list(map(lambda _args: square(*_args), self.test_data))

    def test_non_daemon_nested_workerpool(self):
        """
        Tests nested WorkerPools when daemon==False, which should work
        """
        with WorkerPool(n_jobs=4, daemon=False, start_method='spawn') as pool:
            # Obtain results using nested WorkerPools
            results = pool.map(self._square_daemon, ((X,) for X in repeat(self.test_data, 4)), chunk_size=1)

            # Each of the results should match
            for results_list in results:
                self.assertIsInstance(results_list, list)
                self.assertEqual(self.test_desired_output, results_list)

    def test_daemon_nested_workerpool(self):
        """
        Tests nested WorkerPools when daemon==True, which should not work
        """
        with self.assertRaises(AssertionError), WorkerPool(n_jobs=4, daemon=True, start_method='spawn') as pool:
            pool.map(self._square_daemon, ((X,) for X in repeat(self.test_data, 4)), chunk_size=1)

    @staticmethod
    def _square_daemon(x):
        with WorkerPool(n_jobs=2) as pool:
            return pool.map(square, x, chunk_size=1)


class CPUPinningTest(unittest.TestCase):

    def setUp(self):
        # Create some test data. Note that the regular map reads the inputs as a list of single tuples (one argument),
        # whereas parallel.map sees it as a list of argument lists. Therefore we give the regular map a lambda function
        # which mimics the parallel.map behavior.
        self.test_data = list(enumerate([1, 2, 3, 5, 6, 9, 37, 42, 1337, 0, 3, 5, 0]))
        self.test_desired_output = list(map(lambda _args: square(*_args), self.test_data))

    def test_cpu_pinning(self):
        """
        Test that when parameters are valid, nothing breaks and the pinning is actually happening
        """
        for n_jobs, cpu_ids, expected_mask in [(None, [0], [[0]] * cpu_count()),
                                               (None, [[0, 3]], [[0, 3]] * cpu_count()),
                                               (1, [0], [[0]]),
                                               (1, [[0, 3]], [[0, 3]]),
                                               (2, [0], [[0], [0]]),
                                               (2, [0, 1], [[0], [1]]),
                                               (2, [[0, 3]], [[0, 3], [0, 3]]),
                                               (2, [[0, 1], [0, 1]], [[0, 1], [0, 1]]),
                                               (4, [0], [[0], [0], [0], [0]]),
                                               (4, [0, 1, 2, 3], [[0], [1], [2], [3]]),
                                               (4, [[0, 3]], [[0, 3], [0, 3], [0, 3], [0, 3]])]:
            # The test has been designed for a system with at least 4 cores. We'll skip those test cases where the CPU
            # IDs exceed the number of CPUs.
            if cpu_ids is not None and np.array(cpu_ids).max(initial=0) >= cpu_count():
                continue

            with self.subTest(n_jobs=n_jobs, cpu_ids=cpu_ids), patch('mpire.pool.set_cpu_affinity') as p, \
                    WorkerPool(n_jobs=n_jobs, cpu_ids=cpu_ids) as pool:

                # Verify results
                results_list = pool.map(square, self.test_data)
                self.assertIsInstance(results_list, list)
                self.assertEqual(self.test_desired_output, results_list)

                # Verify that when CPU pinning is used, it is called as many times as there are jobs and is called for
                # each worker process ID
                if cpu_ids is None:
                    self.assertEqual(p.call_args_list, [])
                else:
                    self.assertEqual(p.call_count, pool.pool_params.n_jobs)
                    mask = [call[0][1] for call in p.call_args_list]
                    self.assertListEqual(mask, expected_mask)

    def test_start_methods(self):
        """
        Test for different start methods
        """
        # This test will fail if there are less CPUs available than specified.
        if cpu_count() >= 2:
            n_jobs, cpu_ids, expected_mask = 2, [1, 0], [[1], [0]]
        else:
            n_jobs, cpu_ids, expected_mask = 1, [0], [[0]]

        for start_method in TEST_START_METHODS:
            if start_method == 'threading':
                continue
            with self.subTest(start_method=start_method), patch('mpire.pool.set_cpu_affinity') as p, \
                    WorkerPool(n_jobs=n_jobs, cpu_ids=cpu_ids, start_method=start_method) as pool:
                # Verify results
                results_list = pool.map(square, self.test_data)
                self.assertIsInstance(results_list, list)
                self.assertEqual(self.test_desired_output, results_list)

                # Verify that CPU pinning is used as many times as there are jobs and is called for each worker process
                # ID
                self.assertEqual(p.call_count, pool.pool_params.n_jobs)
                mask = [call[0][1] for call in p.call_args_list]
                self.assertListEqual(mask, expected_mask)

        # This won't work for threading
        with self.assertRaises(AttributeError), WorkerPool(n_jobs=n_jobs, cpu_ids=cpu_ids,
                                                           start_method='threading') as pool:
            pool.map(square, self.test_data)


class ProgressBarTest(unittest.TestCase):

    """
    Print statements in these tests are intentional as it will print multiple progress bars
    """

    def setUp(self):
        # Create some test data. Note that the regular map reads the inputs as a list of single tuples (one argument),
        # whereas parallel.map sees it as a list of argument lists. Therefore we give the regular map a lambda function
        # which mimics the parallel.map behavior.
        self.test_data = list(enumerate([1, 2, 3, 5, 6, 9, 37, 42, 1337, 0, 3, 5, 0]))
        self.test_desired_output = list(map(lambda _args: square(*_args), self.test_data))

        # Numpy test data
        self.test_data_numpy = np.random.rand(100, 2)
        self.test_desired_output_numpy = square_numpy(self.test_data_numpy)
        self.test_data_len_numpy = len(self.test_data_numpy)

        # Get original tqdm lock
        self.original_tqdm_lock = tqdm.get_lock()

    def tearDown(self):
        # The TQDM lock is temporarily changed when using a progress bar in MPIRE, here we check if it is restored
        # correctly afterwards.
        self.assertEqual(tqdm.get_lock(), self.original_tqdm_lock)

    def test_valid_progress_bars_regular_input(self):
        """
        Valid progress bars are either False/True
        """
        print()
        for n_jobs, progress_bar in product([None, 1, 2], [True, False]):

            with self.subTest(n_jobs=n_jobs), WorkerPool(n_jobs=n_jobs) as pool:
                results_list = pool.map(square, self.test_data, progress_bar=progress_bar)
                self.assertIsInstance(results_list, list)
                self.assertEqual(self.test_desired_output, results_list)

    def test_valid_progress_bars_numpy_input(self):
        """
        Test with numpy, as that will change the number of tasks
        """
        print()
        for n_jobs, progress_bar in product([None, 1, 2], [True, False]):

            # Should work just fine
            with self.subTest(n_jobs=n_jobs, progress_bar=progress_bar), WorkerPool(n_jobs=n_jobs) as pool:
                results = pool.map(square_numpy, self.test_data_numpy, progress_bar=progress_bar)
                self.assertIsInstance(results, np.ndarray)
                np.testing.assert_array_equal(results, self.test_desired_output_numpy)

    def test_no_input_data(self):
        """
        Test with empty iterable (this failed before)
        """
        print()
        with WorkerPool() as pool:
            self.assertListEqual(pool.map(square, [], progress_bar=True), [])

    def test_progress_bar_options(self):
        """
        Test different progress bar options. Wrong inputs are tested in test_params
        """
        print()
        for progress_bar_options in [{"unit": "km"}, {"unit": "s", "desc": "I'm a pbar!"}, {"colour": "green"}]:
            with self.subTest(progress_bar_options=progress_bar_options), WorkerPool(n_jobs=2) as pool:
                results = pool.map(square, self.test_data, progress_bar=True, progress_bar_options=progress_bar_options)
                self.assertIsInstance(results, list)
                self.assertEqual(self.test_desired_output, results)

    def test_start_methods(self):
        """
        Test for different start methods
        """
        print()
        for start_method in TEST_START_METHODS:
            with self.subTest(start_method=start_method), WorkerPool(n_jobs=2, start_method=start_method) as pool:
                results_list = pool.map(square, self.test_data, progress_bar=True)
                self.assertIsInstance(results_list, list)
                self.assertEqual(self.test_desired_output, results_list)
                
    def test_progres_bar_styles(self):
        """
        Test different progress bar styles. The std style will give updates by overwriting the previous line, so all
        progress update lines will be there (including the 0% from the start). The rich progress bar updates the 
        widget, so the 0% won't be there. The notebook won't update correctly in a terminal, which is fine. This means
        it won't show the 100% here. Finally, the dashboard style won't give any output.
        """
        print()
        for progress_bar_style, expected_outputs in [
            (None, ["None", "0%", "100%", "13/13"]),
            ('std', ["std", "0%", "100%", "13/13"]),
            ('rich', ["rich", "100%", "13/13"]),
            ('notebook', ["notebook", "0%", "0/13"]),
            ('dashboard', []),
        ]:
            # Some progress bars write to stdout, others to stderr. We'll capture both.
            output = io.StringIO()
            with self.subTest(progress_bar_style=progress_bar_style), redirect_stderr(output), \
                    redirect_stdout(output), WorkerPool(n_jobs=2) as pool:
                results_list = pool.map(square, self.test_data, progress_bar=True,
                                        progress_bar_style=progress_bar_style, 
                                        progress_bar_options={"desc": progress_bar_style or "None"})
                self.assertIsInstance(results_list, list)
                self.assertEqual(self.test_desired_output, results_list)
                
                # Check outputs
                for expected_output in expected_outputs:
                    self.assertIn(expected_output, output.getvalue())
                if not expected_outputs:
                    self.assertEqual(output.getvalue(), '')


class KeepAliveTest(unittest.TestCase):

    """
    In these tests we make use of a barrier. This barrier ensures that we increase the counter for each worker. If it
    wasn't there there's a chance that the first, say 3, workers already performed all the available tasks, while the
    4th worker was still spinning up. In that case the poison pill would be inserted before the fourth worker could even
    start a task and therefore couldn't increase the counter value.
    """

    def setUp(self):
        # Create some test data
        self.test_data = [1, 2, 3, 5, 6, 9, 37, 42, 1337, 0, 3, 5, 0]
        self.test_desired_output_f1 = [x * 2 for x in self.test_data]
        self.test_desired_output_f2 = [x * 3 for x in self.test_data]

    def test_dont_keep_alive(self):
        """
        When keep_alive is set to False it should restart workers between map calls. This means the counter is updated
        each time as well.
        """
        for n_jobs in [1, 3]:
            barrier = Barrier(n_jobs)
            counter = Value('i', 0)
            shared = barrier, counter
            with self.subTest(n_jobs=n_jobs), \
                    WorkerPool(n_jobs=n_jobs, shared_objects=shared, use_worker_state=True, keep_alive=False) as pool:

                self.assertListEqual(pool.map(self._f1, self.test_data, worker_init=self._init1),
                                     self.test_desired_output_f1)
                self.assertEqual(counter.value, n_jobs)
                barrier.reset()

                self.assertListEqual(pool.map(self._f1, self.test_data, worker_init=self._init1),
                                     self.test_desired_output_f1)
                self.assertEqual(counter.value, n_jobs * 2)
                barrier.reset()

                self.assertListEqual(pool.map(self._f1, self.test_data, worker_init=self._init1),
                                     self.test_desired_output_f1)
                self.assertEqual(counter.value, n_jobs * 3)
                barrier.reset()

                self.assertListEqual(pool.map(self._f1, self.test_data, worker_init=self._init1),
                                     self.test_desired_output_f1)
                self.assertEqual(counter.value, n_jobs * 4)

    def test_keep_alive(self):
        """
        When keep_alive is set to True it should reuse existing workers between map calls. This means the counter is
        only updated the first time.
        """
        for n_jobs in [1, 3]:
            barrier = Barrier(n_jobs)
            counter = Value('i', 0)
            shared = barrier, counter
            with self.subTest(n_jobs=n_jobs), \
                    WorkerPool(n_jobs=n_jobs, shared_objects=shared, use_worker_state=True, keep_alive=True) as pool:

                self.assertListEqual(pool.map(self._f1, self.test_data, worker_init=self._init1),
                                     self.test_desired_output_f1)
                self.assertEqual(counter.value, n_jobs)
                barrier.reset()

                self.assertListEqual(list(pool.imap(self._f1, self.test_data, worker_init=self._init1)),
                                     self.test_desired_output_f1)
                self.assertEqual(counter.value, n_jobs)
                barrier.reset()

                self.assertListEqual(pool.map(self._f1, self.test_data, worker_init=self._init1),
                                     self.test_desired_output_f1)
                self.assertEqual(counter.value, n_jobs)

    def test_keep_alive_map_params_change(self):
        """
        When keep_alive is set to True it should reuse existing workers between map calls, even when the called
        function, init or exit functions, or the worker lifespan changes
        """
        for n_jobs in [1, 3]:
            barrier = Barrier(n_jobs)
            counter = Value('i', 0)
            shared = barrier, counter
            with self.subTest(n_jobs=n_jobs), warnings.catch_warnings(), \
                    WorkerPool(n_jobs=n_jobs, shared_objects=shared, use_worker_state=True, keep_alive=True) as pool:

                warnings.simplefilter('ignore')
                self.assertListEqual(pool.map(self._f1, self.test_data, worker_lifespan=100, worker_init=self._init1,
                                              worker_exit=self._exit1),
                                     self.test_desired_output_f1)
                self.assertEqual(counter.value, n_jobs)
                barrier.reset()

                self.assertListEqual(list(pool.imap(self._f2, self.test_data, worker_lifespan=100,
                                                    worker_init=self._init1, worker_exit=self._exit2)),
                                     self.test_desired_output_f2)
                self.assertEqual(counter.value, n_jobs)
                barrier.reset()

                self.assertListEqual(pool.map(self._f2, self.test_data, worker_lifespan=200, worker_init=self._init2,
                                              worker_exit=self._exit1),
                                     self.test_desired_output_f2)
                self.assertEqual(counter.value, n_jobs)
                barrier.reset()

                self.assertListEqual(pool.map(self._f1, self.test_data, worker_lifespan=100, worker_init=self._init1,
                                              worker_exit=None),
                                     self.test_desired_output_f1)
                self.assertEqual(counter.value, n_jobs)

    def test_start_methods(self):
        """
        Test for different start methods
        """
        for start_method in TEST_START_METHODS:
            with self.subTest(start_method=start_method), \
                    WorkerPool(n_jobs=2, use_worker_state=True, keep_alive=True, start_method=start_method) as pool:

                barrier = pool.ctx.Barrier(2)
                counter = pool.ctx.Value('i', 0)
                pool.set_shared_objects((barrier, counter))
                self.assertListEqual(pool.map(self._f1, self.test_data, worker_init=self._init1),
                                     self.test_desired_output_f1)
                self.assertEqual(counter.value, 2)
                barrier.reset()

                self.assertListEqual(list(pool.imap(self._f1, self.test_data, worker_init=self._init1)),
                                     self.test_desired_output_f1)
                self.assertEqual(counter.value, 2)
                barrier.reset()

                self.assertListEqual(pool.map(self._f1, self.test_data, worker_init=self._init1),
                                     self.test_desired_output_f1)
                self.assertEqual(counter.value, 2)

    @staticmethod
    def _init1(_, worker_state):
        worker_state['already_counted'] = False

    @staticmethod
    def _init2(_, worker_state):
        worker_state['already_counted'] = False
        worker_state[4] = 2

    @staticmethod
    def _f1(shared, worker_state, x):
        """
        Function that waits for all workers to spin up and increases the counter by one only once per worker,
        returns x * 2
        """
        barrier, counter = shared
        if not worker_state['already_counted']:
            with counter.get_lock():
                counter.value += 1
            worker_state['already_counted'] = True
            barrier.wait()
        return x * 2

    @staticmethod
    def _f2(shared, worker_state, x):
        """
        Function that waits for all workers to spin up and increases the counter by one only once per worker,
        returns x * 3
        """
        barrier, counter = shared
        if not worker_state['already_counted']:
            with counter.get_lock():
                counter.value += 1
            worker_state['already_counted'] = True
            barrier.wait()
        return x * 3

    @staticmethod
    def _exit1(_, worker_state):
        return worker_state['already_counted']

    @staticmethod
    def _exit2(_, worker_state):
        pass


class ExceptionTest(unittest.TestCase):

    def setUp(self):
        # Create some test data. Note that the regular map reads the inputs as a list of single tuples (one argument),
        # whereas parallel.map sees it as a list of argument lists. Therefore we give the regular map a lambda function
        # which mimics the parallel.map behavior.
        self.test_data = list(enumerate([1, 2, 3, 5, 6, 9, 37, 42, 1337, 0, 3, 5, 0]))
        self.test_desired_output = list(map(lambda _args: square(*_args), self.test_data))
        self.test_data_len = len(self.test_data)

        # Get original tqdm lock
        self.original_tqdm_lock = tqdm.get_lock()

    def tearDown(self):
        # The TQDM lock is temporarily changed when using a progress bar in MPIRE, here we check if it is restored
        # correctly afterwards.
        self.assertEqual(tqdm.get_lock(), self.original_tqdm_lock)

    def test_exceptions(self):
        """
        Tests if MPIRE can handle exceptions well
        """
        # This print statement is intentional as it will print multiple progress bars
        print()
        for n_jobs, n_tasks_max_active, worker_lifespan, progress_bar in [
            (1, None, None, False),
            (3, None, None, False),
            (3, 1, None, False),
            (3, None, 1, False),
            (3, None, None, True),
            (3, 1, None, True),
            (3, None, 1, True),
            (3, 1, 1, True)
        ]:
            print(f"========== {n_jobs}, {n_tasks_max_active}, {worker_lifespan}, {progress_bar} ==========")
            with WorkerPool(n_jobs=n_jobs) as pool:

                # Should work for map like functions
                print("----- square_raises, map -----")
                with self.subTest(n_jobs=n_jobs, n_tasks_max_active=n_tasks_max_active, worker_lifespan=worker_lifespan,
                                  progress_bar=progress_bar, function='square_raises', map='map'), \
                     self.assertRaises(ValueError):
                    pool.map(self._square_raises, self.test_data, max_tasks_active=n_tasks_max_active,
                             worker_lifespan=worker_lifespan, progress_bar=progress_bar)

                # Should work for imap like functions
                print("----- square_raises, imap -----")
                with self.subTest(n_jobs=n_jobs, n_tasks_max_active=n_tasks_max_active, worker_lifespan=worker_lifespan,
                                  progress_bar=progress_bar, function='square_raises', map='imap'), \
                     self.assertRaises(ValueError):
                    list(pool.imap_unordered(self._square_raises, self.test_data, max_tasks_active=n_tasks_max_active,
                                             worker_lifespan=worker_lifespan, progress_bar=progress_bar))

                # Should work for map like functions
                print("----- square_raises_on_idx, map -----")
                with self.subTest(n_jobs=n_jobs, n_tasks_max_active=n_tasks_max_active, worker_lifespan=worker_lifespan,
                                  progress_bar=progress_bar, function='square_raises_on_idx', map='map'), \
                     self.assertRaises(ValueError):
                    pool.map(self._square_raises_on_idx, self.test_data, max_tasks_active=n_tasks_max_active,
                             worker_lifespan=worker_lifespan, progress_bar=progress_bar)

                # Should work for imap like functions
                print("----- square_raises_on_idx, imap -----")
                with self.subTest(n_jobs=n_jobs, n_tasks_max_active=n_tasks_max_active, worker_lifespan=worker_lifespan,
                                  progress_bar=progress_bar, function='square_raises_on_idx', map='imap'), \
                     self.assertRaises(ValueError):
                    list(pool.imap_unordered(self._square_raises_on_idx, self.test_data,
                                             max_tasks_active=n_tasks_max_active, worker_lifespan=worker_lifespan,
                                             progress_bar=progress_bar))

    def test_start_methods(self):
        """
        Test for different start methods
        """
        print()
        for start_method, progress_bar in product(TEST_START_METHODS, [False, True]):
            print(f"========== {start_method}, {progress_bar} ==========")
            if RUNNING_WINDOWS and progress_bar and start_method == 'threading':
                print("Not yet supported on Windows")
                continue
            with self.subTest(start_method=start_method, progress_bar=progress_bar), \
                    WorkerPool(n_jobs=2, start_method=start_method) as pool:

                # Should work for map like functions
                print("----- square_raises, map -----")
                with self.subTest(function='square_raises', map='map'), self.assertRaises(ValueError):
                    pool.map(self._square_raises, self.test_data, progress_bar=progress_bar)

                # Should work for imap like functions
                print("----- square_raises, imap -----")
                with self.subTest(function='square_raises', map='imap'), self.assertRaises(ValueError):
                    list(pool.imap_unordered(self._square_raises, self.test_data, progress_bar=progress_bar))

                if not progress_bar:
                    # Should work for apply like functions
                    print("----- square_raises, apply -----")
                    with self.subTest(function='square_raises', func='apply'), self.assertRaises(ValueError):
                        pool.apply(self._square_raises, self.test_data[0])

                # Should work for map like functions
                print("----- square_raises_on_idx, map -----")
                with self.subTest(function='square_raises_on_idx', map='map'), self.assertRaises(ValueError):
                    pool.map(self._square_raises_on_idx, self.test_data, progress_bar=progress_bar)

                # Should work for imap like functions
                print("----- square_raises_on_idx, imap -----")
                with self.subTest(function='square_raises_on_idx', map='imap'), self.assertRaises(ValueError):
                    list(pool.imap_unordered(self._square_raises_on_idx, self.test_data, progress_bar=progress_bar))

    def test_defunct_processes_exit(self):
        """
        Tests if MPIRE correctly shuts down after process becomes defunct using exit()
        """
        print()
        for n_jobs, progress_bar, worker_lifespan in [(1, False, None),
                                                      (3, True, 1),
                                                      (3, False, 3)]:
            for start_method in TEST_START_METHODS:
                # Progress bar on Windows + threading is not supported right now
                if RUNNING_WINDOWS and start_method == 'threading' and progress_bar:
                    continue
                print(f"========== {start_method}, {n_jobs}, {progress_bar}, {worker_lifespan} ==========")
                with self.subTest(n_jobs=n_jobs, progress_bar=progress_bar, worker_lifespan=worker_lifespan,
                                  start_method=start_method), self.assertRaises(SystemExit), \
                        WorkerPool(n_jobs=n_jobs, start_method=start_method) as pool:
                    pool.map(self._exit, range(100), progress_bar=progress_bar, worker_lifespan=worker_lifespan)

    def test_defunct_processes_kill(self):
        """
        Tests if MPIRE correctly shuts down after one process becomes defunct using os.kill().

        We kill worker 0 and to be sure it's alive we set an event object and then go in an infinite loop. The kill
        thread waits until the event is set and then kills the worker. The other workers are also ensured to have done
        something so we can test what happens during restarts
        """
        print()
        for n_jobs, progress_bar, worker_lifespan in [(1, False, None),
                                                      (3, True, 1),
                                                      (3, False, 3)]:
            for start_method in TEST_START_METHODS:
                # Can't kill threads
                if start_method == 'threading':
                    continue

                print(f"========== {start_method}, {n_jobs}, {progress_bar}, {worker_lifespan} ==========")
                with self.subTest(n_jobs=n_jobs, progress_bar=progress_bar, worker_lifespan=worker_lifespan,
                                  start_method=start_method), self.assertRaises(RuntimeError), \
                        WorkerPool(n_jobs=n_jobs, pass_worker_id=True, start_method=start_method) as pool:
                    events = [pool.ctx.Event() for _ in range(n_jobs)]
                    kill_thread = Thread(target=self._kill_process, args=(events[0], pool))
                    kill_thread.start()
                    pool.set_shared_objects(events)
                    pool.map(self._worker_0_sleeps_others_square, range(100), progress_bar=progress_bar,
                             worker_lifespan=worker_lifespan, chunk_size=1)

    def test_dill_deadlock(self):
        """
        Exceptions on the queue need to be flushed before the worker is terminated. This is one example where it used
        to cause a deadlock (https://github.com/Slimmer-AI/mpire/issues/56)
        """
        data = [(x, y, z) for x, y, z in zip(range(0, 100), range(42, 142), range(10, -90, -1))]
        with self.assertRaises(ZeroDivisionError), WorkerPool(n_jobs=5, use_dill=True) as pool:
            for _ in pool.imap(lambda x, y, z: x * y / z, data):
                pass

    @staticmethod
    def _square_raises(_, x):
        raise ValueError(x)

    @staticmethod
    def _square_raises_on_idx(idx, x):
        if idx == 5:
            raise ValueError(x)
        else:
            return idx, x * x

    @staticmethod
    def _exit(_):
        exit()

    @staticmethod
    def _worker_0_sleeps_others_square(worker_id, events, x):
        """
        Worker 0 waits until the other workers have at least spun up and then sets her event and sleeps
        """
        if worker_id == 0:
            [event.wait() for event in events[1:]]
            events[0].set()
            while True:
                pass
        else:
            events[worker_id].set()
            return x * x

    @staticmethod
    def _kill_process(event, pool):
        """
        Wait for event and kill
        """
        event.wait()
        pool._workers[0].terminate()


class TimeoutTest(unittest.TestCase):

    def setUp(self):
        # Create some test data
        self.test_data = [1, 2, 3]

    def test_worker_init_timeout(self):
        """
        Checks if the worker_init timeout is properly triggered
        """
        print()
        for start_method in TEST_START_METHODS:

            print(f"========== {start_method}, well below timeout ==========")
            with self.subTest('Well below timeout', start_method=start_method), \
                    WorkerPool(2, start_method=start_method) as pool:
                self.assertListEqual(pool.map(self._f1, self.test_data, worker_init=self._init1,
                                              worker_init_timeout=100), self.test_data)

            print(f"========== {start_method}, exceeding timeout, map ==========")
            with self.subTest('Exceeding timeout, map', start_method=start_method), \
                    WorkerPool(2, start_method=start_method) as pool, self.assertRaises(TimeoutError):
                pool.map(self._f1, self.test_data, worker_init=self._init2, worker_init_timeout=0.01)

            print(f"========== {start_method}, exceeding timeout, imap ==========")
            with self.subTest('Exceeding timeout, imap', start_method=start_method), \
                    WorkerPool(2, start_method=start_method) as pool, self.assertRaises(TimeoutError):
                for _ in pool.imap(self._f1, self.test_data, worker_init=self._init2, worker_init_timeout=0.01):
                    pass

            print(f"========== {start_method}, exceeding timeout, apply ==========")
            with self.subTest('Exceeding timeout, apply', start_method=start_method), \
                    WorkerPool(2, start_method=start_method) as pool, self.assertRaises(TimeoutError):
                pool.apply(self._f1, self.test_data[0], worker_init=self._init2, worker_init_timeout=0.01)

    def test_worker_task_timeout(self):
        """
        Checks if the worker_init timeout is properly triggered
        """
        print()
        for start_method in TEST_START_METHODS:

            print(f"========== {start_method}, well below timeout ==========")
            with self.subTest('Well below timeout', start_method=start_method), \
                    WorkerPool(2, start_method=start_method) as pool:
                self.assertListEqual(pool.map(self._f1, self.test_data, task_timeout=100), self.test_data)

            print(f"========== {start_method}, exceeding timeout, map ==========")
            with self.subTest('Exceeding timeout, map', start_method=start_method), \
                    WorkerPool(2, start_method=start_method) as pool, self.assertRaises(TimeoutError):
                pool.map(self._f2, self.test_data, task_timeout=0.01)

            print(f"========== {start_method}, exceeding timeout, imap ==========")
            with self.subTest('Exceeding timeout, imap', start_method=start_method), \
                    WorkerPool(2, start_method=start_method) as pool, self.assertRaises(TimeoutError):
                for _ in pool.imap(self._f2, self.test_data, task_timeout=0.01):
                    pass

            print(f"========== {start_method}, exceeding timeout, apply ==========")
            with self.subTest('Exceeding timeout, apply', start_method=start_method), \
                    WorkerPool(2, start_method=start_method) as pool, self.assertRaises(TimeoutError):
                pool.apply(self._f2, self.test_data[0], task_timeout=0.01)

    def test_worker_exit_timeout(self):
        """
        Checks if the worker_exit timeout is properly triggered
        """
        print()
        for start_method in TEST_START_METHODS:

            print(f"========== {start_method}, well below timeout ==========")
            with self.subTest('Well below timeout', start_method=start_method), \
                    WorkerPool(2, start_method=start_method) as pool:
                self.assertListEqual(pool.map(self._f1, self.test_data, worker_exit=self._exit1,
                                              worker_exit_timeout=100), self.test_data)

            print(f"========== {start_method}, exceeding timeout, map ==========")
            with self.subTest('Exceeding timeout, map', start_method=start_method), \
                    WorkerPool(2, start_method=start_method) as pool, self.assertRaises(TimeoutError):
                pool.map(self._f1, self.test_data, worker_exit=self._exit2, worker_exit_timeout=0.01)

            print(f"========== {start_method}, exceeding timeout, imap ==========")
            with self.subTest('Exceeding timeout, imap', start_method=start_method), \
                    WorkerPool(2, start_method=start_method) as pool, self.assertRaises(TimeoutError):
                for _ in pool.imap(self._f1, self.test_data, worker_exit=self._exit2, worker_exit_timeout=0.01):
                    pass

            print(f"========== {start_method}, exceeding timeout, apply ==========")
            with self.subTest('Exceeding timeout, apply', start_method=start_method), \
                    WorkerPool(2, start_method=start_method) as pool, self.assertRaises(TimeoutError):
                pool.apply(self._f1, self.test_data[0], worker_exit=self._exit2, worker_exit_timeout=0.01)
                pool.stop_and_join()

    def test_apply_async_multiple_task_timeout(self):
        """ Test that some apply_async() tasks time out correctly and don't kill the whole pool """
        print()
        for start_method in tqdm(TEST_START_METHODS):
            with WorkerPool(n_jobs=3, start_method=start_method) as pool:
                results = [pool.apply_async(self._f3, (i,), task_timeout=0.1) for i in range(6)]
                for i, result in enumerate(results):
                    if i % 2 == 0:
                        self.assertEqual(result.get(), i)
                    else:
                        with self.assertRaises(TimeoutError):
                            result.get()

    @staticmethod
    def _init1():
        pass

    @staticmethod
    def _init2():
        time.sleep(1)

    @staticmethod
    def _f1(x):
        return x

    @staticmethod
    def _f2(x):
        time.sleep(1)
        return x

    @staticmethod
    def _exit1():
        pass

    @staticmethod
    def _exit2():
        time.sleep(1)

    @staticmethod
    def _f3(x):
        if x % 2 == 0:
            return x
        else:
            time.sleep(1)
            return x


class OrderTasksTest(unittest.TestCase):
    """
    Tests if the tasks are properly ordered
    """

    def setUp(self):
        # Create some test data
        self.test_data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]

    def test_order_tasks(self):
        """
        Checks if the tasks are properly ordered
        """
        for start_method in TEST_START_METHODS:
            with self.subTest(start_method=start_method, chunk_size=1), \
                    WorkerPool(4, start_method=start_method, pass_worker_id=True, use_worker_state=True,
                               order_tasks=True) as pool:
                pool.map_unordered(self._f, self.test_data, worker_init=self._init, worker_exit=self._exit,
                                   chunk_size=1)
                exit_results = sorted(pool.get_exit_results(), key=lambda state: state['worker_id'])
                self.assertListEqual(exit_results, [{'worker_id': 0, 'tasks': [1, 5, 9, 13, 17]},
                                                    {'worker_id': 1, 'tasks': [2, 6, 10, 14, 18]},
                                                    {'worker_id': 2, 'tasks': [3, 7, 11, 15, 19]},
                                                    {'worker_id': 3, 'tasks': [4, 8, 12, 16, 20]}])

            with self.subTest(start_method=start_method, chunk_size=3), \
                    WorkerPool(4, start_method=start_method, pass_worker_id=True, use_worker_state=True,
                               order_tasks=True) as pool:
                pool.map_unordered(self._f, self.test_data, worker_init=self._init, worker_exit=self._exit,
                                   chunk_size=3)
                exit_results = sorted(pool.get_exit_results(), key=lambda state: state['worker_id'])
                self.assertListEqual(exit_results, [{'worker_id': 0, 'tasks': [1, 2, 3, 13, 14, 15]},
                                                    {'worker_id': 1, 'tasks': [4, 5, 6, 16, 17, 18]},
                                                    {'worker_id': 2, 'tasks': [7, 8, 9, 19, 20]},
                                                    {'worker_id': 3, 'tasks': [10, 11, 12]}])

    def test_order_tasks_twice(self):
        """
        Checks if the tasks are properly ordered the second time around as well.
        """
        for start_method in TEST_START_METHODS:
            with self.subTest(start_method=start_method, chunk_size=3, keep_alive=True), \
                    WorkerPool(4, start_method=start_method, pass_worker_id=True, use_worker_state=True,
                               order_tasks=True, keep_alive=True) as pool:
                pool.map_unordered(self._f, self.test_data, worker_init=self._init, worker_exit=self._exit,
                                   chunk_size=3)
                pool.map_unordered(self._f, self.test_data, worker_init=self._init, worker_exit=self._exit,
                                   chunk_size=3)
                pool.stop_and_join()
                exit_results = sorted(pool.get_exit_results(), key=lambda state: state['worker_id'])
                self.assertListEqual(exit_results,
                                     [{'worker_id': 0, 'tasks': [1, 2, 3, 13, 14, 15, 1, 2, 3, 13, 14, 15]},
                                      {'worker_id': 1, 'tasks': [4, 5, 6, 16, 17, 18, 4, 5, 6, 16, 17, 18]},
                                      {'worker_id': 2, 'tasks': [7, 8, 9, 19, 20, 7, 8, 9, 19, 20]},
                                      {'worker_id': 3, 'tasks': [10, 11, 12, 10, 11, 12]}])

    @staticmethod
    def _init(wid, state):
        state['worker_id'] = wid
        state['tasks'] = []

    @staticmethod
    def _f(_, state, x):
        state['tasks'].append(x)

    @staticmethod
    def _exit(_, state):
        return state