Testing a module with input() and print() inside - using TestCase

Situation:

A students task is to write a python script, doing something with only input() and print().

Example task: "Write python script that asks for number 'n' and prints string of 'n' stars."

He writes a script like this:

solution.py

n = int(input())
print('*' * n)

I want to write python module, that tests that his script is working correctly.

My solution till now is to do this:

test.py from io import StringIO import sys

sys.stdout = StringIO()
sys.stdin = StringIO(str(2))

import riesenie

s = sys.stdout.getvalue()
if s != '**' * 2 + '\n':
    raise Exception('WRONG')
else:
    print("OK", file=sys.stderr)

BUT, I would like to achieve this behaviour using TestCase. Something like:

test.py from unittest import TestCase

class TestSolution(TestCase):
    def test_stars(self):
        sys.stdout = StringIO()
        sys.stdin = StringIO(str(2))
        import riesenie
        s = sys.stdout.getvalue()

        self.assertEqual('**', s)

But this is not working.

Is there a way to achieve this? Thank you in advance.