"""Tests for timeout."""

from __future__ import absolute_import
import signal
import time
import unittest

from . import timeout

LONG_TIME_SECS = 10


class TimeoutTest(unittest.TestCase):

  def testRunWithoutTimeout(self):
    """Tests that handler is not called if there is no timeout."""
    def _Handler(signum, frame):
      del signum, frame  # Unused.
      self.fail('Timeout should not occur.')

    with timeout.timeout(handler=_Handler, timeout_secs=LONG_TIME_SECS):
      pass

  def testRunWithZeroTimeout(self):
    """Tests that SIGALRM isn't raised with a 0 second timeout."""
    def _Handler(signum, frame):
      del signum, frame  # Unused.
      self.fail('Timeout should not occur.')

    with timeout.timeout(handler=_Handler, timeout_secs=0):
      pass

  def testDefaultHandler(self):
    """Tests the default handler raises the correct exception."""
    with self.assertRaises(timeout.TimeoutException):
      with timeout.timeout(timeout_secs=1):
        time.sleep(LONG_TIME_SECS)

  def testRunWithTimeout(self):
    """Tests timing out with a custom timeout handler."""

    class TestException(Exception):
      """A custom exception for this test."""

    def _Handler(signum, frame):
      del frame  # Unused.
      self.assertEqual(signum, signal.SIGALRM)
      raise TestException()

    with self.assertRaises(TestException):
      with timeout.timeout(handler=_Handler, timeout_secs=1):
        time.sleep(LONG_TIME_SECS)

  def testPreviouslySetHandlerIsCalledWithoutTimeout(self):
    """Tests that previous handler is called if there is no timeout."""

    class TestException(Exception):
      """A custom exception for this test."""

    def _FirstHandler(signum, frame):
      del frame  # Unused.
      self.assertEqual(signum, signal.SIGALRM)
      raise TestException()

    def _SecondHandler(signum, frame):
      del signum, frame  # Unused.
      self.fail('Timeout should not occur.')

    with self.assertRaises(TestException):
      signal.signal(signal.SIGALRM, _FirstHandler)
      signal.alarm(2)
      with timeout.timeout(handler=_SecondHandler, timeout_secs=LONG_TIME_SECS):
        pass
      time.sleep(LONG_TIME_SECS)

  def testPreviouslySetHandlerIsCalledWithTimeout(self):
    """Tests that previous handler is called if there is a timeout."""

    class TestException(Exception):
      """A custom exception for this test."""

    def _FirstHandler(signum, frame):
      del frame  # Unused.
      self.assertEqual(signum, signal.SIGALRM)
      raise TestException()

    def _SecondHandler(signum, frame):
      del frame  # Unused.
      self.assertEqual(signum, signal.SIGALRM)
      raise TestException()

    with self.assertRaises(TestException):
      signal.signal(signal.SIGALRM, _FirstHandler)
      signal.alarm(2)
      with self.assertRaises(TestException):
        with timeout.timeout(handler=_SecondHandler, timeout_secs=1):
          time.sleep(LONG_TIME_SECS)
      time.sleep(LONG_TIME_SECS)


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