Ben Heckmann
Ben Heckmann

Reputation: 43

How to mock input and output for testing (for competitive programming)?

Let's say I have a function that prints an integer read from the console:

in submission.py:

def print_input():
    n = int(input())
    print(n)

How can I patch the input and output to test this function on an array of test cases? (Obviously, I would like the mock methods for input/output to behave exactly like input() and print() in python.

Upvotes: 0

Views: 1116

Answers (2)

ukBaz
ukBaz

Reputation: 8014

You can use unittest.mock to give the return value from input and capture the sys.stdout from the print statement. e.g:

import unittest
from unittest import mock
import io


def print_input():
    n = int(input())
    print(n)


class Test101(unittest.TestCase):
    def test_name(self):
        with mock.patch('sys.stdout', new=io.StringIO()) as fake_out:
            with mock.patch('builtins.input', return_value="1"):
                print_input()
                self.assertEqual("1\n", fake_out.getvalue())


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

Upvotes: 1

AlanSTACK
AlanSTACK

Reputation: 6065

The easiest way is to just define a new input(...) function on top of your code

def input():
    """
    monkey patch method that will override default
    """
    return 10

...

[your code here]

Upvotes: 0

Related Questions