#!/usr/bin/env python

#

# Unit tests for the multiprocessing package

#

import unittest

import threading

import Queue

import time

import sys

import os

import gc

import signal

import array

import copy

import socket

import random

import logging

# Work around broken sem_open implementations

try:

    import multiprocessing.synchronize

except ImportError, e:

    from test.test_support import TestSkipped

    raise TestSkipped(e)

import multiprocessing.dummy

import multiprocessing.connection

import multiprocessing.managers

import multiprocessing.heap

import multiprocessing.pool

import _multiprocessing

from multiprocessing import util

#

#

#

latin = str

#

# Constants

#

LOG_LEVEL = util.SUBWARNING

#LOG_LEVEL = logging.WARNING

DELTA = 0.1

CHECK_TIMINGS = False     # making true makes tests take a lot longer

                          # and can sometimes cause some non-serious

                          # failures because some calls block a bit

                          # longer than expected

if CHECK_TIMINGS:

    TIMEOUT1, TIMEOUT2, TIMEOUT3 = 0.82, 0.35, 1.4

else:

    TIMEOUT1, TIMEOUT2, TIMEOUT3 = 0.1, 0.1, 0.1

HAVE_GETVALUE = not getattr(_multiprocessing,

                            'HAVE_BROKEN_SEM_GETVALUE', False)

#

# Creates a wrapper for a function which records the time it takes to finish

#

class TimingWrapper(object):

    def __init__(self, func):

        self.func = func

        self.elapsed = None

    def __call__(self, *args, **kwds):

        t = time.time()

        try:

            return self.func(*args, **kwds)

        finally:

            self.elapsed = time.time() - t

#

# Base class for test cases

#

class BaseTestCase(object):

    ALLOWED_TYPES = ('processes', 'manager', 'threads')

    def assertTimingAlmostEqual(self, a, b):

        if CHECK_TIMINGS:

            self.assertAlmostEqual(a, b, 1)

    def assertReturnsIfImplemented(self, value, func, *args):

        try:

            res = func(*args)

        except NotImplementedError:

            pass

        else:

            return self.assertEqual(value, res)

#

# Return the value of a semaphore

#

def get_value(self):

    try:

        return self.get_value()

    except AttributeError:

        try:

            return self._Semaphore__value

        except AttributeError:

            try:

                return self._value

            except AttributeError:

                raise NotImplementedError

#

# Testcases

#

class _TestProcess(BaseTestCase):

    ALLOWED_TYPES = ('processes', 'threads')

    def test_current(self):

        if self.TYPE == 'threads':

            return

        current = self.current_process()

        authkey = current.authkey

        self.assertTrue(current.is_alive())

        self.assertTrue(not current.daemon)

        self.assertTrue(isinstance(authkey, bytes))

        self.assertTrue(len(authkey) > 0)

        self.assertEqual(current.ident, os.getpid())

        self.assertEqual(current.exitcode, None)

    def _test(self, q, *args, **kwds):

        current = self.current_process()

        q.put(args)

        q.put(kwds)

        q.put(current.name)

        if self.TYPE != 'threads':

            q.put(bytes(current.authkey))

            q.put(current.pid)

    def test_process(self):

        q = self.Queue(1)

        e = self.Event()

        args = (q, 1, 2)

        kwargs = {'hello':23, 'bye':2.54}

        name = 'SomeProcess'

        p = self.Process(

            target=self._test, args=args, kwargs=kwargs, name=name

            )

        p.daemon = True

        current = self.current_process()

        if self.TYPE != 'threads':

            self.assertEquals(p.authkey, current.authkey)

        self.assertEquals(p.is_alive(), False)

        self.assertEquals(p.daemon, True)

        self.assertTrue(p not in self.active_children())

        self.assertTrue(type(self.active_children()) is list)

        self.assertEqual(p.exitcode, None)

        p.start()

        self.assertEquals(p.exitcode, None)

        self.assertEquals(p.is_alive(), True)

        self.assertTrue(p in self.active_children())

        self.assertEquals(q.get(), args[1:])

        self.assertEquals(q.get(), kwargs)

        self.assertEquals(q.get(), p.name)

        if self.TYPE != 'threads':

            self.assertEquals(q.get(), current.authkey)

            self.assertEquals(q.get(), p.pid)

        p.join()

        self.assertEquals(p.exitcode, 0)

        self.assertEquals(p.is_alive(), False)

        self.assertTrue(p not in self.active_children())

    def _test_terminate(self):

        time.sleep(1000)

    def test_terminate(self):

        if self.TYPE == 'threads':

            return

        p = self.Process(target=self._test_terminate)

        p.daemon = True

        p.start()

        self.assertEqual(p.is_alive(), True)

        self.assertTrue(p in self.active_children())

        self.assertEqual(p.exitcode, None)

        p.terminate()

        join = TimingWrapper(p.join)

        self.assertEqual(join(), None)

        self.assertTimingAlmostEqual(join.elapsed, 0.0)

        self.assertEqual(p.is_alive(), False)

        self.assertTrue(p not in self.active_children())

        p.join()

        # XXX sometimes get p.exitcode == 0 on Windows ...

        #self.assertEqual(p.exitcode, -signal.SIGTERM)

    def test_cpu_count(self):

        try:

            cpus = multiprocessing.cpu_count()

        except NotImplementedError:

            cpus = 1

        self.assertTrue(type(cpus) is int)

        self.assertTrue(cpus >= 1)

    def test_active_children(self):

        self.assertEqual(type(self.active_children()), list)

        p = self.Process(target=time.sleep, args=(DELTA,))

        self.assertTrue(p not in self.active_children())

        p.start()

        self.assertTrue(p in self.active_children())

        p.join()

        self.assertTrue(p not in self.active_children())

    def _test_recursion(self, wconn, id):

        from multiprocessing import forking

        wconn.send(id)

        if len(id) < 2:

            for i in range(2):

                p = self.Process(

                    target=self._test_recursion, args=(wconn, id+[i])

                    )

                p.start()

                p.join()

    def test_recursion(self):

        rconn, wconn = self.Pipe(duplex=False)

        self._test_recursion(wconn, [])

        time.sleep(DELTA)

        result = []

        while rconn.poll():

            result.append(rconn.recv())

        expected = [

            [],

              [0],

                [0, 0],

                [0, 1],

              [1],

                [1, 0],

                [1, 1]

            ]

        self.assertEqual(result, expected)

#

#

#

class _UpperCaser(multiprocessing.Process):

    def __init__(self):

        multiprocessing.Process.__init__(self)

        self.child_conn, self.parent_conn = multiprocessing.Pipe()

    def run(self):

        self.parent_conn.close()

        for s in iter(self.child_conn.recv, None):

            self.child_conn.send(s.upper())

        self.child_conn.close()

    def submit(self, s):

        assert type(s) is str

        self.parent_conn.send(s)

        return self.parent_conn.recv()

    def stop(self):

        self.parent_conn.send(None)

        self.parent_conn.close()

        self.child_conn.close()

class _TestSubclassingProcess(BaseTestCase):

    ALLOWED_TYPES = ('processes',)

    def test_subclassing(self):

        uppercaser = _UpperCaser()

        uppercaser.start()

        self.assertEqual(uppercaser.submit('hello'), 'HELLO')

        self.assertEqual(uppercaser.submit('world'), 'WORLD')

        uppercaser.stop()

        uppercaser.join()

#

#

#

def queue_empty(q):

    if hasattr(q, 'empty'):

        return q.empty()

    else:

        return q.qsize() == 0

def queue_full(q, maxsize):

    if hasattr(q, 'full'):

        return q.full()

    else:

        return q.qsize() == maxsize

class _TestQueue(BaseTestCase):

    def _test_put(self, queue, child_can_start, parent_can_continue):

        child_can_start.wait()

        for i in range(6):

            queue.get()

        parent_can_continue.set()

    def test_put(self):

        MAXSIZE = 6

        queue = self.Queue(maxsize=MAXSIZE)

        child_can_start = self.Event()

        parent_can_continue = self.Event()

        proc = self.Process(

            target=self._test_put,

            args=(queue, child_can_start, parent_can_continue)

            )

        proc.daemon = True

        proc.start()

        self.assertEqual(queue_empty(queue), True)

        self.assertEqual(queue_full(queue, MAXSIZE), False)

        queue.put(1)

        queue.put(2, True)

        queue.put(3, True, None)

        queue.put(4, False)

        queue.put(5, False, None)

        queue.put_nowait(6)

        # the values may be in buffer but not yet in pipe so sleep a bit

        time.sleep(DELTA)

        self.assertEqual(queue_empty(queue), False)

        self.assertEqual(queue_full(queue, MAXSIZE), True)

        put = TimingWrapper(queue.put)

        put_nowait = TimingWrapper(queue.put_nowait)

        self.assertRaises(Queue.Full, put, 7, False)

        self.assertTimingAlmostEqual(put.elapsed, 0)

        self.assertRaises(Queue.Full, put, 7, False, None)

        self.assertTimingAlmostEqual(put.elapsed, 0)

        self.assertRaises(Queue.Full, put_nowait, 7)

        self.assertTimingAlmostEqual(put_nowait.elapsed, 0)

        self.assertRaises(Queue.Full, put, 7, True, TIMEOUT1)

        self.assertTimingAlmostEqual(put.elapsed, TIMEOUT1)

        self.assertRaises(Queue.Full, put, 7, False, TIMEOUT2)

        self.assertTimingAlmostEqual(put.elapsed, 0)

        self.assertRaises(Queue.Full, put, 7, True, timeout=TIMEOUT3)

        self.assertTimingAlmostEqual(put.elapsed, TIMEOUT3)

        child_can_start.set()

        parent_can_continue.wait()

        self.assertEqual(queue_empty(queue), True)

        self.assertEqual(queue_full(queue, MAXSIZE), False)

        proc.join()

    def _test_get(self, queue, child_can_start, parent_can_continue):

        child_can_start.wait()

        #queue.put(1)

        queue.put(2)

        queue.put(3)

        queue.put(4)

        queue.put(5)

        parent_can_continue.set()

    def test_get(self):

        queue = self.Queue()

        child_can_start = self.Event()

        parent_can_continue = self.Event()

        proc = self.Process(

            target=self._test_get,

            args=(queue, child_can_start, parent_can_continue)

            )

        proc.daemon = True

        proc.start()

        self.assertEqual(queue_empty(queue), True)

        child_can_start.set()

        parent_can_continue.wait()

        time.sleep(DELTA)

        self.assertEqual(queue_empty(queue), False)

        # Hangs unexpectedly, remove for now

        #self.assertEqual(queue.get(), 1)

        self.assertEqual(queue.get(True, None), 2)

        self.assertEqual(queue.get(True), 3)

        self.assertEqual(queue.get(timeout=1), 4)

        self.assertEqual(queue.get_nowait(), 5)

        self.assertEqual(queue_empty(queue), True)

        get = TimingWrapper(queue.get)

        get_nowait = TimingWrapper(queue.get_nowait)

        self.assertRaises(Queue.Empty, get, False)

        self.assertTimingAlmostEqual(get.elapsed, 0)

        self.assertRaises(Queue.Empty, get, False, None)

        self.assertTimingAlmostEqual(get.elapsed, 0)

        self.assertRaises(Queue.Empty, get_nowait)

        self.assertTimingAlmostEqual(get_nowait.elapsed, 0)

        self.assertRaises(Queue.Empty, get, True, TIMEOUT1)

        self.assertTimingAlmostEqual(get.elapsed, TIMEOUT1)

        self.assertRaises(Queue.Empty, get, False, TIMEOUT2)

        self.assertTimingAlmostEqual(get.elapsed, 0)

        self.assertRaises(Queue.Empty, get, timeout=TIMEOUT3)

        self.assertTimingAlmostEqual(get.elapsed, TIMEOUT3)

        proc.join()

    def _test_fork(self, queue):

        for i in range(10, 20):

            queue.put(i)

        # note that at this point the items may only be buffered, so the

        # process cannot shutdown until the feeder thread has finished

        # pushing items onto the pipe.

    def test_fork(self):

        # Old versions of Queue would fail to create a new feeder

        # thread for a forked process if the original process had its

        # own feeder thread.  This test checks that this no longer

        # happens.

        queue = self.Queue()

        # put items on queue so that main process starts a feeder thread

        for i in range(10):

            queue.put(i)

        # wait to make sure thread starts before we fork a new process

        time.sleep(DELTA)

        # fork process

        p = self.Process(target=self._test_fork, args=(queue,))

        p.start()

        # check that all expected items are in the queue

        for i in range(20):

            self.assertEqual(queue.get(), i)

        self.assertRaises(Queue.Empty, queue.get, False)

        p.join()

    def test_qsize(self):

        q = self.Queue()

        try:

            self.assertEqual(q.qsize(), 0)

        except NotImplementedError:

            return

        q.put(1)

        self.assertEqual(q.qsize(), 1)

        q.put(5)

        self.assertEqual(q.qsize(), 2)

        q.get()

        self.assertEqual(q.qsize(), 1)

        q.get()

        self.assertEqual(q.qsize(), 0)

    def _test_task_done(self, q):

        for obj in iter(q.get, None):

            time.sleep(DELTA)

            q.task_done()

    def test_task_done(self):

        queue = self.JoinableQueue()

        if sys.version_info < (2, 5) and not hasattr(queue, 'task_done'):

            return

        workers = [self.Process(target=self._test_task_done, args=(queue,))

                   for i in xrange(4)]

        for p in workers:

            p.start()

        for i in xrange(10):

            queue.put(i)

        queue.join()

        for p in workers:

            queue.put(None)

        for p in workers:

            p.join()

#

#

#