import unittest
import io
from unittest.mock import patch
from hanoi import hanoi, prompt_for_integer

class TestHanoi(unittest.TestCase):
   
   def test_hanoi_zero_disks(self):
        expected_output = "No disks, nothing to do\n"
        # with self.assertRaises(SystemExit):
        with unittest.mock.patch('sys.stdout', new=io.StringIO()) as mock_stdout:
           hanoi(0, 'A', 'B', 'C')
   
   def test_hanoi_one_disk(self):
       expected_output = "Move disk 1 from A to C\n"
       with unittest.mock.patch('sys.stdout', new=io.StringIO()) as mock_stdout:
           hanoi(1, 'A', 'B', 'C')
           self.assertEqual(mock_stdout.getvalue(), expected_output)
   
   def test_hanoi_two_disks(self):
       expected_output = "Move disk 1 from A to B\nMove disk 2 from A to C\nMove disk 1 from B to C\n"
       with unittest.mock.patch('sys.stdout', new=io.StringIO()) as mock_stdout:
           hanoi(2, 'A', 'B', 'C')
           self.assertEqual(mock_stdout.getvalue(), expected_output)
   
   def test_hanoi_three_disks(self):
       expected_output = "Move disk 1 from A to C\nMove disk 2 from A to B\nMove disk 1 from C to B\nMove disk 3 from A to C\nMove disk 1 from B to A\nMove disk 2 from B to C\nMove disk 1 from A to C\n"
       with unittest.mock.patch('sys.stdout', new=io.StringIO()) as mock_stdout:
           hanoi(3, 'A', 'B', 'C')
           self.assertEqual(mock_stdout.getvalue(), expected_output)
   
   def test_hanoi_negative_input(self):
        expected_output = "No disks, nothing to do\n"
        # with self.assertRaises(SystemExit):
        with unittest.mock.patch('sys.stdout', new=io.StringIO()) as mock_stdout:
            hanoi(-1, 'A', 'B', 'C')
   
   def test_prompt_for_integer_valid_integer(self):
       with unittest.mock.patch('builtins.input', return_value='4'):
            self.assertEqual(prompt_for_integer(), 4)
   
   def test_prompt_for_integer_invalid_input(self):
       with unittest.mock.patch('builtins.input', side_effect=['abc', '5']):
           self.assertEqual(prompt_for_integer(), 5)

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