嵌套if语句的Python单元测试

Question

所以我一直在发布单元测试问题，因为我试图让他们很好。我会尽量做到尽可能清楚。下面我有一个嵌套的if语句，我想模拟difficulty变量的输入，并检查stdout是否符合条件。

这里是我的运行代码..

def main(): 
    print("''''''''''''''''''''''''''''''''''''") 
    print("''''''''''''''''''''''''''''''''''''") 
    print("''' \t\t\t\t '''") 
    print("''' \t\t\t\t '''") 
    print("''' \t Pirate Games\t\t '''") 
    print("''' \tStart a new game?\t '''") 
    print("''' \t\t\t\t '''") 
    print("''' \t\t\t\t '''") 
    print("''''''''''''''''''''''''''''''''''''") 
    print("''''''''''''''''''''''''''''''''''''") 

    newGame = input("").lower() 

    if newGame == "yes" or "y": 

     print("1.Scallywag\n2.Crew\n3.Pirate") 
     difficulty = input("Choose ye toughness!") 

     if difficulty == "1": 
      print("TEST") 

    elif newGame == "no" or "n": 
     print("Goodbye! Come Again") 

    else: 
     print("Enter correct input!") 

现在，这里是我的单元测试，但这只是对于第一个if语句，它工作正常。我不知道如何去做第二个。任何帮助表示赞赏，谢谢大家提前。

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

from RunFile import main 

class GetInputTest(unittest.TestCase): 

    @patch('builtins.input', return_value='yes') 
    def test_output(self,m): 
     saved_stdout = sys.stdout 
     try: 
      out = io.StringIO() 
      sys.stdout = out 
      main() 
      output = out.getvalue().strip() 
      self.assertEqual(output, "1.Scallywag\n2.Crew\n3.Pirate") 
     finally: 
      sys.stdout = saved_stdout 


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

Answer 1

可以使用side_effect参数patch使input回报"yes"在第一时间其所谓，和"1"第二次：如果你正在运行3.4

class GetInputTest(unittest.TestCase): 

    @patch('builtins.input', side_effect=["yes", "1"]) 
    def test_output(self,m): 
     saved_stdout = sys.stdout 
     try: 
      out = io.StringIO() 
      sys.stdout = out 
      main() 
      output = out.getvalue().strip() 
      # Make sure TEST appears at the end, in addition to the original list of items. 
      self.assertEqual(output, "1.Scallywag\n2.Crew\n3.Pirate\nTEST") 

     finally: 
      sys.stdout = saved_stdout