Python unit test for nested if statement

问题

So I've been posting unit test questions because I'm trying to get good at them. I'll try to be as clear as possible. Below I have a nested if statement, I want to mock the input for the difficulty variable and check for the stdout if the condition is met.

here is my run code..

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!")

Now here's my unit test, but this is only for the first If Statement, which works fine. I don't know how to go about doing the second one. Any help is appreciated, thank you all in advance.

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

回答1:

You can use the side_effect argument to patch to make input return "yes" the first time its called, and "1" the second time:

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

来源：https://stackoverflow.com/questions/25677260/python-unit-test-for-nested-if-statement

标签

python

unit-testing

if-statement

nested