Skip to Content

How to write the UnitTest with stdin at Pytest

If you want to write UnitTest when using stdin in Python. Pytest provide setattr function in monkeypatch

from io import StringIO
import sys


def divide():
    input = sys.stdin.readline
    return list(input())


def gather():
    input = sys.stdin.readline
    return sum(list(map(int, input().split())))


def test_divide(monkeypatch):
    monkeypatch.setattr('sys.stdin', StringIO('abc'))
    assert divide() == ['a', 'b', 'c']


def test_gather(monkeypatch):
    monkeypatch.setattr('sys.stdin', StringIO('1 2 3'))
    assert gather() == 6

Reference

See Also