Search code examples
pythonpython-3.xunit-testingpython-unittest

Unit-test for python3 script using unittest


I have a little python3 script. Something like this

import sys


content = sys.stdin.read()

print(content)

I need to write a unit test for this using import unittest. I am trying to understand, how simulate stdin using unittest.

I found some articles here

and here, but I still couldn't deal with it, because after I run testing scripts, program expects me to give stdin.

Could you help me to rewrite this code (or write another) for testing my script

import unittest
from unittest.mock import patch

import module_under_test


class MyTestCase(unittest.TestCase):

    def setUp(self):
        # raw_input is untouched before test
        assert module_under_test.raw_input is __builtins__.raw_input

    def test_using_with(self):
        input_data = "123"
        expected = int(input_data)

        with patch.object(module_under_test, "raw_input", create=True, 
                return_value=expected):
            # create=True is needed as raw_input is not in the globals of 
            # module_under_test, but actually found in __builtins__ .
            actual = module_under_test.function()

        self.assertEqual(expected, actual)

    @patch.object(module_under_test, "raw_input", create=True)
    def test_using_decorator(self, raw_input):
        raw_input.return_value = input_data = "123"
        expected = int(input_data)

        actual = module_under_test.function()

        self.assertEqual(expected, actual)

    def tearDown(self):
        # raw input is restored after test
        assert module_under_test.raw_input is __builtins__.raw_input

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

Solution

  • You should only need a single test method that uses patch:

    script.py:

    import sys
    
    
    def main():
      content = sys.stdin.read()
      print(content)
    
    
    if __name__ == '__main__':
      main()
    
    

    script_test.py:

    import unittest
    from unittest.mock import patch
    
    import script
    
    
    class ScriptTest(unittest.TestCase):
    
      @patch('sys.stdin.read')
      def test_content_reading(self, mock_stdin):
        mock_stdin.return_value = 'mocked data 123'
        with patch('builtins.print') as mock_print:
          script.main()
          mock_print.assert_called_once_with('mocked data 123')
    
    
    if __name__ == '__main__':
      unittest.main()
    
    

    Example Usage:

    $ python -m unittest script_test.py
    .
    ----------------------------------------------------------------------
    Ran 1 test in 0.032s
    
    OK
    

    NB. When using sys.stdin.read(). It waits for EOF (End-Of-File) to read the content. If you're testing it in a terminal, you can signify EOF by pressing Ctrl-D on Linux/macOS or Ctrl-Z on Windows.