嵌套 if 语句的 Python 单元测试

Question

所以我一直在发布单元测试问题，因为我正在努力解决这些问题。我会尽量说清楚。下面我有一个嵌套的 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()

Answer 1

您可以使用patch 的side_effect 参数使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