mitogen/tests/io_op_test.py

import errno
import select

try:
    from unittest import mock
except ImportError:
    import mock

import testlib
import mitogen.core


def py24_mock_fix(m):
    def wrapper(*args, **kwargs):
        ret = m(*args, **kwargs)
        if isinstance(ret, Exception):
            raise ret
        return ret
    return wrapper


class RestartTest(object):
    func = staticmethod(mitogen.core.io_op)
    exception_class = None

    def test_eintr_restarts(self):
        m = mock.Mock()
        m.side_effect = [
            self.exception_class(errno.EINTR),
            self.exception_class(errno.EINTR),
            self.exception_class(errno.EINTR),
            'yay',
        ]
        rc, disconnected = self.func(py24_mock_fix(m), 'input')
        self.assertEqual(rc, 'yay')
        self.assertFalse(disconnected)
        self.assertEqual(4, m.call_count)
        self.assertEqual(m.mock_calls, [
            mock.call('input'),
            mock.call('input'),
            mock.call('input'),
            mock.call('input'),
        ])


class SelectRestartTest(RestartTest, testlib.TestCase):
    exception_class = select.error


class OsErrorRestartTest(RestartTest, testlib.TestCase):
    exception_class = OSError


class DisconnectTest(object):
    func = staticmethod(mitogen.core.io_op)
    errno = None
    exception_class = None

    def test_disconnection(self):
        m = mock.Mock()
        m.side_effect = self.exception_class(self.errno)
        rc, disconnected = self.func(m, 'input')
        self.assertEqual(rc, None)
        self.assertTrue(disconnected)
        self.assertEqual(1, m.call_count)
        self.assertEqual(m.mock_calls, [
            mock.call('input'),
        ])


class SelectDisconnectEioTest(DisconnectTest, testlib.TestCase):
    errno = errno.EIO
    exception_class = select.error


class SelectDisconnectEconnresetTest(DisconnectTest, testlib.TestCase):
    errno = errno.ECONNRESET
    exception_class = select.error


class SelectDisconnectEpipeTest(DisconnectTest, testlib.TestCase):
    errno = errno.EPIPE
    exception_class = select.error


class OsErrorDisconnectEioTest(DisconnectTest, testlib.TestCase):
    errno = errno.EIO
    exception_class = OSError


class OsErrorDisconnectEconnresetTest(DisconnectTest, testlib.TestCase):
    errno = errno.ECONNRESET
    exception_class = OSError


class OsErrorDisconnectEpipeTest(DisconnectTest, testlib.TestCase):
    errno = errno.EPIPE
    exception_class = OSError


class ExceptionTest(object):
    func = staticmethod(mitogen.core.io_op)
    errno = None
    exception_class = None

    def test_exception(self):
        m = mock.Mock()
        m.side_effect = self.exception_class(self.errno)
        e = self.assertRaises(self.exception_class,
                              lambda: self.func(m, 'input'))
        self.assertEqual(e, m.side_effect)
        self.assertEqual(1, m.call_count)
        self.assertEqual(m.mock_calls, [
            mock.call('input'),
        ])


class SelectExceptionTest(ExceptionTest, testlib.TestCase):
    errno = errno.EBADF
    exception_class = select.error


class OsErrorExceptionTest(ExceptionTest, testlib.TestCase):
    errno = errno.EBADF
    exception_class = OSError