1# Copyright 2024 The Chromium Authors 2# Use of this source code is governed by a BSD-style license that can be 3# found in the LICENSE file. 4 5import enum 6import unittest 7 8from crossbench.helper.state import (BaseState, StateMachine, 9 UnexpectedStateError) 10from tests import test_helper 11 12 13@enum.unique 14class CustomState(BaseState): 15 INITIAL = enum.auto() 16 READY = enum.auto() 17 DONE = enum.auto() 18 19 20class StateMachineTestCase(unittest.TestCase): 21 22 def test_init(self): 23 state_machine = StateMachine(CustomState.INITIAL) 24 self.assertIs(state_machine.state, CustomState.INITIAL) 25 state_machine = StateMachine(CustomState.READY) 26 self.assertIs(state_machine.state, CustomState.READY) 27 28 def test_eq(self): 29 state_machine = StateMachine(CustomState.READY) 30 state_machine_2 = StateMachine(CustomState.READY) 31 self.assertEqual(state_machine, state_machine) 32 self.assertEqual(state_machine, state_machine_2) 33 self.assertEqual(state_machine, CustomState.READY) 34 self.assertNotEqual(state_machine, None) 35 self.assertNotEqual(state_machine, CustomState.INITIAL) 36 self.assertNotEqual(state_machine, StateMachine(CustomState.INITIAL)) 37 38 def test_transition(self): 39 state_machine = StateMachine(CustomState.INITIAL) 40 state_machine.transition(CustomState.INITIAL, to=CustomState.READY) 41 self.assertEqual(state_machine.state, CustomState.READY) 42 with self.assertRaises(UnexpectedStateError) as cm: 43 state_machine.transition(CustomState.INITIAL, to=CustomState.READY) 44 self.assertIn("INITIAL", str(cm.exception)) 45 self.assertIn("READY", str(cm.exception)) 46 47 def test_transition_multi_current(self): 48 state_machine = StateMachine(CustomState.INITIAL) 49 state_machine.transition( 50 CustomState.INITIAL, CustomState.READY, to=CustomState.READY) 51 self.assertEqual(state_machine.state, CustomState.READY) 52 state_machine.transition( 53 CustomState.INITIAL, CustomState.READY, to=CustomState.READY) 54 self.assertEqual(state_machine.state, CustomState.READY) 55 state_machine.transition( 56 CustomState.INITIAL, CustomState.READY, to=CustomState.DONE) 57 self.assertEqual(state_machine.state, CustomState.DONE) 58 with self.assertRaises(UnexpectedStateError) as cm: 59 state_machine.transition( 60 CustomState.INITIAL, CustomState.READY, to=CustomState.DONE) 61 self.assertIn("INITIAL", str(cm.exception)) 62 self.assertIn("READY", str(cm.exception)) 63 self.assertIn("DONE", str(cm.exception)) 64 65 def test_expect(self): 66 state_machine = StateMachine(CustomState.INITIAL) 67 state_machine.expect(CustomState.INITIAL) 68 with self.assertRaises(RuntimeError) as cm: 69 state_machine.expect(CustomState.READY) 70 self.assertIn("INITIAL", str(cm.exception)) 71 self.assertIn("READY", str(cm.exception)) 72 73 def test_expect_before(self): 74 state_machine = StateMachine(CustomState.INITIAL) 75 state_machine.expect_before(CustomState.READY) 76 state_machine.expect_before(CustomState.DONE) 77 78 state_machine.transition(CustomState.INITIAL, to=CustomState.READY) 79 with self.assertRaises(UnexpectedStateError) as cm: 80 state_machine.expect_before(CustomState.READY) 81 self.assertEqual(cm.exception.state, CustomState.READY) 82 self.assertEqual(cm.exception.expected, (CustomState.INITIAL,)) 83 state_machine.expect_before(CustomState.DONE) 84 85 state_machine.transition(CustomState.READY, to=CustomState.DONE) 86 with self.assertRaises(UnexpectedStateError) as cm: 87 state_machine.expect_before(CustomState.DONE) 88 self.assertEqual(cm.exception.state, CustomState.DONE) 89 self.assertEqual(cm.exception.expected, 90 (CustomState.INITIAL, CustomState.READY)) 91 92 def test_expect_at_least(self): 93 state_machine = StateMachine(CustomState.INITIAL) 94 with self.assertRaises(UnexpectedStateError) as cm: 95 state_machine.expect_at_least(CustomState.READY) 96 self.assertEqual(cm.exception.state, CustomState.INITIAL) 97 self.assertEqual(cm.exception.expected, 98 (CustomState.READY, CustomState.DONE)) 99 with self.assertRaises(UnexpectedStateError) as cm: 100 state_machine.expect_at_least(CustomState.DONE) 101 self.assertEqual(cm.exception.state, CustomState.INITIAL) 102 self.assertEqual(cm.exception.expected, (CustomState.DONE,)) 103 104 state_machine.transition(CustomState.INITIAL, to=CustomState.READY) 105 state_machine.expect_at_least(CustomState.INITIAL) 106 state_machine.expect_at_least(CustomState.READY) 107 with self.assertRaises(UnexpectedStateError) as cm: 108 state_machine.expect_at_least(CustomState.DONE) 109 self.assertEqual(cm.exception.state, CustomState.READY) 110 self.assertEqual(cm.exception.expected, (CustomState.DONE,)) 111 112 state_machine.transition(CustomState.READY, to=CustomState.DONE) 113 state_machine.expect_at_least(CustomState.INITIAL) 114 state_machine.expect_at_least(CustomState.READY) 115 state_machine.expect_at_least(CustomState.DONE) 116 117 118if __name__ == "__main__": 119 test_helper.run_pytest(__file__) 120