How can I mock multiple prints?

Question

I have this function:

def function1():
    number1 = int(input())
    number2 = int(input())
    print("Calculate sum")
    sum = number1 + number2
    print("Calculate product")
    product = number1 * number2
    print("Calculate substraction")
    sub = number1 - number2
    print("Done")

I have to do an unit test for this function but I don't know how to check all the prints. I implemented this function that can verify the last print but I want to verify with assert and mock all the printed messages.

@patch('sys.stdout', new_callable=StringIO)
def test_main(mock_stdout):
    input_main = [6, 2]
    with patch('builtins.input', side_effect=input_main):
        main()
        assert mock_stdout.getvalue() == 'Done\n'

test_main()

Answer 1

You can use the native pytest capability instead :

def test_myoutput(capsys):  # or use "capfd" for fd-level
    print("hello")
    sys.stderr.write("world\n")
    captured = capsys.readouterr()
    assert captured.out == "hello\n"
    assert captured.err == "world\n"
    print("next")
    captured = capsys.readouterr()
    assert captured.out == "next\n"

Answer 2

Using unittest you can do:

import unittest
from unittest.mock import patch, call

@patch('builtins.print')
def test_your_function(self, mock_print):
   your_function()
   self.assertEqual(mock_print.mock_calls, [call("Calculate sum"), 
                                            call("Calculate product"),
                                            call("Calculate substraction"),
                                            call("Done")])