How to test a python file that contains no functions and has multiple input calls?

Question

I have the following python file two_strings.py:

str1 = input()
str2 = input()
combined = str1 + "," + str2
print(combined)

I then have a jupyter notebook file called test_two_strings.ipynb:

import sys
import unittest
from unittest.mock import patch
import two_strings
from io import StringIO

sys.path.append('Challenge10/')

class SimpleTestCase(unittest.TestCase):
    def test_defs(self):    
        defs = [v for v in dir(two_strings) if v[0] != '_']
        assert 'str1' in defs, f'Definition for str1 not found'
        assert 'str2' in defs, f'Definition for str2 not found'
        assert 'combined' in defs, f'Definition for combined not found'
    
    def test_inputs_outputs(self):
        user_input = [
            ['a','b'],
            ['10','sixty six'],
            ['hello','there']
            ['once upon a time','in fairyland']
        ]
        expected_output = [
            'a,b',
            '10,sixty six',
            'hello,there',
            'once upon a time,in fairyland'
        ]
        for i in range(len(user_input)):
            with patch('builtins.input', side_effect=user_input[i]):
                output = StringIO()
                with patch('sys.stdout', new=output):
                    two_strings.main()
                    self.assertEqual(output.getvalue().strip(), expected_output[i])
        
    
if __name__ == '__main__':
    unittest.main(argv=['-v'],verbosity=2, exit=False)

That will test the python file within the notebook file. I am using jupyterlite. However I am unable to figure out exactly how to setup unit tests for testing multiple lines of input() calls and making sure the output is correct. I know the expected output is correct but knowing how janky using input in jupyter notebooks are in jupyterlite, I want to know if this is possible? If so, how?

Answer 1

After many days of testing and looking through documentation it is actually quite simple, as pythons sys module allows individuals to directly interact with stdout and stdin:

import sys
from io import StringIO
import unittest

challengeFile = 'challenge.py'
fileRunner = open(challengeFile)
code = fileRunner.read()

class TestName(unittest.TestCase):
    def get_stdout(self, inputs):
        original_stdin = sys.stdin
        original_stdout = sys.stdout
        test_inputs = inputs
        sys.stdin = StringIO('\n'.join(test_inputs))
        sys.stdout = StringIO()
        exec(code)
        script_output = sys.stdout.getvalue()
        sys.stdin = original_stdin
        sys.stdout = original_stdout
        return script_output
    
    def correct(self, test_output, expected_output):
        self.assertEqual(str(test_output).strip(), str(expected_output).strip())
    
    def test_challenge_1(self):
        self.correct(self.get_stdout(["1 1 1 1 1 1"]), "6")

    def test_challenge_2(self):
        self.correct(self.get_stdout(["1 5"]), "6")

    def test_challenge_3(self):
        self.correct(self.get_stdout(["8 11"]), "19")
        
    fileRunner.close()
    
unittest.main(exit=False)

This will test the following code (notice how there are no functions!):

s = input().split()
sList = list(map(lambda x: int(x), s))
print(sum(sList))