Testing printed output in Python 2 and Python 3 - mikec964/chelmbigstock GitHub Wiki

What if you have a method that prints a result, and you want to make sure that the printed output is correct? In Python 2 and Python 3, the solutions are very different. The general idea in both, though, is to remap stdout into a variable where it can be tested.

Python 2

Here's the relevant code (check the source for the full code)::

import StringIO

class test_poi(unittest.TestCase):
    """Tests for the Poi object"""

    def test_poi_printing(self):
        fake_out = StringIO.StringIO()
        sys.stdout = fake_out
        poi_list = Poi.load_file('testdata-poi.txt')
        poi_list[0].print_csv()
        sys.stdout = sys.__stdout__
        self.assertEqual(fake_out.getvalue(), 'Tiberium,26,1500,(704, 271)\n')

        fake_out = StringIO.StringIO()
        sys.stdout = fake_out
        poi_list[0].print_table()
        sys.stdout = sys.__stdout__
        self.assertEqual(fake_out.getvalue(), 'Tiberium   26  1500  (704, 271)\n')

Note that you have to assign the fake_out variable for each test.

Python 3

    def test_poi_printing(self):
        with mock.patch('sys.stdout', new=StringIO()) as fake_out:
            poi_list[0].print_csv()
            self.assertEqual(fake_out.getvalue(), 'Tiberium,26,1500,(704, 271)')
            poi_list[0].print_table()
            self.assertEqual(fake_out.getvalue(), 'Tiberium   26  1500  (704, 271)')

Note that mock is built into Python 3. It exists as a library for Python 2.7, but I wasn't able to make it work like this.