mitogen/tests/timer_test.py

import time

import mock
import unittest2

import mitogen.core
import mitogen.parent

import testlib


class TimerListMixin(object):
    klass = mitogen.parent.TimerList

    def setUp(self):
        self.list = self.klass()


class GetTimeoutTest(TimerListMixin, testlib.TestCase):
    def test_empty(self):
        self.assertEquals(None, self.list.get_timeout())

    def test_one_event(self):
        self.list.schedule(2, lambda: None)
        self.list._now = lambda: 1
        self.assertEquals(1, self.list.get_timeout())

    def test_two_events_same_moment(self):
        self.list.schedule(2, lambda: None)
        self.list.schedule(2, lambda: None)
        self.list._now = lambda: 1
        self.assertEquals(1, self.list.get_timeout())

    def test_two_events(self):
        self.list.schedule(2, lambda: None)
        self.list.schedule(3, lambda: None)
        self.list._now = lambda: 1
        self.assertEquals(1, self.list.get_timeout())

    def test_two_events_expired(self):
        self.list.schedule(2, lambda: None)
        self.list.schedule(3, lambda: None)
        self.list._now = lambda: 3
        self.assertEquals(0, self.list.get_timeout())

    def test_two_events_in_past(self):
        self.list.schedule(2, lambda: None)
        self.list.schedule(3, lambda: None)
        self.list._now = lambda: 30
        self.assertEquals(0, self.list.get_timeout())

    def test_two_events_in_past(self):
        self.list.schedule(2, lambda: None)
        self.list.schedule(3, lambda: None)
        self.list._now = lambda: 30
        self.assertEquals(0, self.list.get_timeout())

    def test_one_cancelled(self):
        t1 = self.list.schedule(2, lambda: None)
        t2 = self.list.schedule(3, lambda: None)
        self.list._now = lambda: 0
        t1.cancel()
        self.assertEquals(3, self.list.get_timeout())

    def test_two_cancelled(self):
        t1 = self.list.schedule(2, lambda: None)
        t2 = self.list.schedule(3, lambda: None)
        self.list._now = lambda: 0
        t1.cancel()
        t2.cancel()
        self.assertEquals(None, self.list.get_timeout())


class ScheduleTest(TimerListMixin, testlib.TestCase):
    def test_in_past(self):
        self.list._now = lambda: 30
        timer = self.list.schedule(29, lambda: None)
        self.assertEquals(29, timer.when)
        self.assertEquals(0, self.list.get_timeout())

    def test_in_future(self):
        self.list._now = lambda: 30
        timer = self.list.schedule(31, lambda: None)
        self.assertEquals(31, timer.when)
        self.assertEquals(1, self.list.get_timeout())

    def test_same_moment(self):
        self.list._now = lambda: 30
        timer = self.list.schedule(31, lambda: None)
        timer2 = self.list.schedule(31, lambda: None)
        self.assertEquals(31, timer.when)
        self.assertEquals(31, timer2.when)
        self.assertTrue(timer is not timer2)
        self.assertEquals(1, self.list.get_timeout())


class ExpireTest(TimerListMixin, testlib.TestCase):
    def test_in_past(self):
        timer = self.list.schedule(29, mock.Mock())
        self.list._now = lambda: 30
        self.list.expire()
        self.assertEquals(1, len(timer.func.mock_calls))

    def test_in_future(self):
        timer = self.list.schedule(29, mock.Mock())
        self.list._now = lambda: 28
        self.list.expire()
        self.assertEquals(0, len(timer.func.mock_calls))

    def test_same_moment(self):
        timer = self.list.schedule(29, mock.Mock())
        timer2 = self.list.schedule(29, mock.Mock())
        self.list._now = lambda: 29
        self.list.expire()
        self.assertEquals(1, len(timer.func.mock_calls))
        self.assertEquals(1, len(timer2.func.mock_calls))

    def test_cancelled(self):
        self.list._now = lambda: 29
        timer = self.list.schedule(29, mock.Mock())
        timer.cancel()
        self.assertEquals(None, self.list.get_timeout())
        self.list._now = lambda: 29
        self.list.expire()
        self.assertEquals(0, len(timer.func.mock_calls))
        self.assertEquals(None, self.list.get_timeout())


class CancelTest(TimerListMixin, testlib.TestCase):
    def test_single_cancel(self):
        self.list._now = lambda: 29
        timer = self.list.schedule(29, mock.Mock())
        timer.cancel()
        self.list.expire()
        self.assertEquals(0, len(timer.func.mock_calls))

    def test_double_cancel(self):
        self.list._now = lambda: 29
        timer = self.list.schedule(29, mock.Mock())
        timer.cancel()
        timer.cancel()
        self.list.expire()
        self.assertEquals(0, len(timer.func.mock_calls))


@mitogen.core.takes_econtext
def do_timer_test_econtext(econtext):
    do_timer_test(econtext.broker)


def do_timer_test(broker):
    now = time.time()
    latch = mitogen.core.Latch()
    broker.defer(lambda:
        broker.timers.schedule(
            now + 0.250,
            lambda: latch.put('hi'),
        )
    )

    assert 'hi' == latch.get()
    assert time.time() > (now + 0.250)


class BrokerTimerTest(testlib.TestCase):
    klass = mitogen.master.Broker

    def test_call_later(self):
        broker = self.klass()
        try:
            do_timer_test(broker)
        finally:
            broker.shutdown()
            broker.join()

    def test_child_upgrade(self):
        router = mitogen.master.Router()
        try:
            c = router.local()
            c.call(mitogen.parent.upgrade_router)
            c.call(do_timer_test_econtext)
        finally:
            router.broker.shutdown()
            router.broker.join()


if __name__ == '__main__':
    unittest2.main()