import builtins
import unittest


def main():
    return int(input('Enter number: '))


class TestMethods(unittest.TestCase):
    def setUp(self):
        # let's store original function somewhere
        self.original_input = builtins.input

    def test(self):
        # we are replacing built-in function input() with our own function
        builtins.input = lambda x: '2'
        self.assertEqual(main(), 2)

    def tearDown(self):
        # let's restore original function
        builtins.input = self.original_input


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