import sys
import unittest
import warnings
import linecache

class BaseTest(unittest.TestCase):
    four_string = 'c:4: Warning: a\n'
    five_string = 'c:5: Warning: b\n  e\n'

    def setUp(self):
        super(BaseTest, self).setUp()
        self._format = warnings.formatwarning
        self._show = warnings.showwarning
        self._checkcache = linecache.checkcache

    def tearDown(self):
        super(BaseTest, self).tearDown()
        warnings.formatwarning = self._format
        warnings.showwarning = self._show
        linecache.checkcache = self._checkcache

    def test_four(self):
        try:
            self.assertEqual(warnings.formatwarning('a', Warning, 'c', 4),
                             self.four_string)
        finally:
            warnings.formatwarning = self._format
            warnings.showwarning = self._show

    def test_five(self):
        try:
            self.assertEqual(warnings.formatwarning('b', Warning, 'c', 5,
                                                    '  e  '), self.five_string)
        finally:
            warnings.formatwarning = self._format
            warnings.showwarning = self._show


IDLE_FOUR = ('\nWarning (from warnings module):\n'
             '  File "c", line 4\nWarning: a\n>>> ')
IDLE_FIVE = ('\nWarning (from warnings module):\n'
                   '  File "c", line 5\n    e\nWarning: b\n>>> ')

class WithRun(BaseTest):
    four_string = IDLE_FOUR[:-4]
    five_string = IDLE_FIVE[:-4]

    def setUp(self):
        super(WithRun, self).setUp()
        import idlelib.run

    def tearDown(self):
        del sys.modules['idlelib.run']
        super(WithRun, self).tearDown()



class WithShell(BaseTest):
    four_string = IDLE_FOUR
    five_string = IDLE_FIVE

    def setUp(self):
        super(WithShell, self).setUp()
        import idlelib.PyShell

    def tearDown(self):
        del sys.modules['idlelib.PyShell']
        super(WithShell, self).tearDown()



class WithShellRun(WithShell):
    four_string = IDLE_FOUR[: -4]
    five_string = IDLE_FIVE[: -4]

    def setUp(self):
        super(WithShellRun, self).setUp()
        import idlelib.run

    def tearDown(self):
        del sys.modules['idlelib.run']
        super(WithShellRun, self).tearDown()



class WithRunShell(BaseTest):
    four_string = IDLE_FOUR
    five_string = IDLE_FIVE

    def setUp(self):
        super(WithRunShell, self).setUp()
        import idlelib.PyShell

    def tearDown(self):
        del sys.modules['idlelib.PyShell']
        super(WithRunShell, self).tearDown()



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