嵌套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