嵌套if语句的Python单元测试

所以我一直在发布单元测试问题，因为我正在努力做好这些问题。我会尽量说清楚。下面我有一个嵌套的 if 语句，我想模拟 difficulty 变量的输入并检查标准输出是否满足条件。

这是我的运行代码..

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() 

请您参考如下方法:

您可以使用 side_effect patch 的参数，使 input 在第一次调用时返回 "yes"，在第二次调用时返回 "1"时间:

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