How can I mock multiple prints?

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()