Python Unit test to capture print output

Question

I would like some help writing a unit test that would test a print output. Here's my code:

import random

friends_list = ["Monica","Joey", "Pheobe", "Chandler", "Rachel", "Ross"]

Friend=random.choice(friends_list)

print(friend)

How would I test to ensure the "Friend" printed?

So far I tried:

def test_friends():
friends.

Not sure what to do after this

Answer 1

You can use the unittest module to write a unit test to ensure that the name of the friend is printed.

import unittest
from io import StringIO
import sys
import random

def get_print_output():
    captured_output = StringIO()
    sys.stdout = captured_output
    friends_list = ["Monica", "Joey", "Pheobe", "Chandler", "Rachel", "Ross"]
    Friend = random.choice(friends_list)
    print(Friend)
    sys.stdout = sys.__stdout__
    return captured_output.getvalue().strip()

# Your test class
class TestPrintFriend(unittest.TestCase):

    # Test method to check if the printed name is in the list of friends
    def test_printed_friend(self):
        printed_output = get_print_output()
        friends_list = ["Monica", "Joey", "Pheobe", "Chandler", "Rachel", "Ross"]
        self.assertIn(printed_output, friends_list)

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

Output:

.
----------------------------------------------------------------------
Ran 1 test in 0.001s

OK