import unittest
from unittest.mock import MagicMock

from analogai.foundation.common.design_patterns.state_machine import StateMachine

class TestStateMachine(unittest.TestCase):
    def setUp(self):
        self.context = MagicMock()
        self.state_machine = StateMachine(self.context)
        
        self.execute_result = MagicMock()
        self.mock_state_instance = MagicMock()
        self.mock_state_instance.execute.return_value = self.execute_result
        
        self.mock_state_class = MagicMock()
    
    def test_initialization(self):
        self.assertEqual(self.state_machine.context, self.context)
        self.assertIsNone(self.state_machine.current_state)
        self.assertIsNone(self.state_machine.previous_state)
    
    def test_state_to_first_transition(self):
        self.state_machine.state_to(self.mock_state_class)
        
        self.assertEqual(self.state_machine.current_state, self.mock_state_class)
        self.assertIsNone(self.state_machine.previous_state)
    
    def test_state_to_subsequent_transition(self):
        first_state = MagicMock()
        self.state_machine.current_state = first_state
        
        self.state_machine.state_to(self.mock_state_class)
        
        self.assertEqual(self.state_machine.current_state, self.mock_state_class)
        self.assertEqual(self.state_machine.previous_state, first_state)
    
    def test_set_context(self):
        new_context = MagicMock()
        self.state_machine.set_context(new_context)
        self.assertEqual(self.state_machine.context, new_context)

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