1from unittest import TestCase, skipIf
2
3try:
4    import enum
5except ImportError:
6    enum = None
7
8from transitions.extensions import MachineFactory
9from .test_pygraphviz import pgv
10from .test_graphviz import pgv as gv
11
12
13@skipIf(enum is None, "enum is not available")
14class TestEnumsAsStates(TestCase):
15
16    def setUp(self):
17        class States(enum.Enum):
18            RED = 1
19            YELLOW = 2
20            GREEN = 3
21        self.machine_cls = MachineFactory.get_predefined()
22        self.States = States
23
24    def test_pass_enums_as_states(self):
25        m = self.machine_cls(states=self.States, initial=self.States.YELLOW)
26
27        assert m.state == self.States.YELLOW
28        assert m.is_RED() is False
29        assert m.is_YELLOW() is True
30        assert m.is_RED() is False
31
32        m.to_RED()
33
34        assert m.state == self.States.RED
35        assert m.is_RED() is True
36        assert m.is_YELLOW() is False
37        assert m.is_GREEN() is False
38
39    def test_transitions(self):
40        m = self.machine_cls(states=self.States, initial=self.States.RED)
41        m.add_transition('switch_to_yellow', self.States.RED, self.States.YELLOW)
42        m.add_transition('switch_to_green', 'YELLOW', 'GREEN')
43
44        m.switch_to_yellow()
45        assert m.is_YELLOW() is True
46
47        m.switch_to_green()
48        assert m.is_YELLOW() is False
49        assert m.is_GREEN() is True
50
51    def test_if_enum_has_string_behavior(self):
52        class States(str, enum.Enum):
53            __metaclass__ = enum.EnumMeta
54
55            RED = 'red'
56            YELLOW = 'yellow'
57
58        m = self.machine_cls(states=States, auto_transitions=False, initial=States.RED)
59        m.add_transition('switch_to_yellow', States.RED, States.YELLOW)
60
61        m.switch_to_yellow()
62        assert m.is_YELLOW() is True
63
64    def test_property_initial(self):
65        transitions = [
66            {'trigger': 'switch_to_yellow', 'source': self.States.RED, 'dest': self.States.YELLOW},
67            {'trigger': 'switch_to_green', 'source': 'YELLOW', 'dest': 'GREEN'},
68        ]
69
70        m = self.machine_cls(states=self.States, initial=self.States.RED, transitions=transitions)
71        m.switch_to_yellow()
72        assert m.is_YELLOW()
73
74        m.switch_to_green()
75        assert m.is_GREEN()
76
77    def test_pass_state_instances_instead_of_names(self):
78        state_A = self.machine_cls.state_cls(self.States.YELLOW)
79        state_B = self.machine_cls.state_cls(self.States.GREEN)
80
81        states = [state_A, state_B]
82
83        m = self.machine_cls(states=states, initial=state_A)
84        assert m.state == self.States.YELLOW
85
86        m.add_transition('advance', state_A, state_B)
87        m.advance()
88        assert m.state == self.States.GREEN
89
90    def test_state_change_listeners(self):
91        class States(enum.Enum):
92            ONE = 1
93            TWO = 2
94
95        class Stuff(object):
96            def __init__(self, machine_cls):
97                self.state = None
98                self.machine = machine_cls(states=States, initial=States.ONE, model=self)
99
100                self.machine.add_transition('advance', States.ONE, States.TWO)
101                self.machine.add_transition('reverse', States.TWO, States.ONE)
102                self.machine.on_enter_TWO('hello')
103                self.machine.on_exit_TWO('goodbye')
104
105            def hello(self):
106                self.message = 'Hello'
107
108            def goodbye(self):
109                self.message = 'Goodbye'
110
111        s = Stuff(self.machine_cls)
112        s.advance()
113
114        assert s.is_TWO()
115        assert s.message == 'Hello'
116
117        s.reverse()
118
119        assert s.is_ONE()
120        assert s.message == 'Goodbye'
121
122    def test_enum_zero(self):
123        from enum import IntEnum
124
125        class State(IntEnum):
126            FOO = 0
127            BAR = 1
128
129        transitions = [
130            ['foo', State.FOO, State.BAR],
131            ['bar', State.BAR, State.FOO]
132        ]
133
134        m = self.machine_cls(states=State, initial=State.FOO, transitions=transitions)
135        m.foo()
136        self.assertTrue(m.is_BAR())
137        m.bar()
138        self.assertTrue(m.is_FOO())
139
140    def test_get_transitions(self):
141        m = self.machine_cls(states=self.States, initial=self.States.RED)
142        self.assertEqual(3, len(m.get_transitions(source=self.States.RED)))
143        self.assertEqual(3, len(m.get_transitions(dest=self.States.RED)))
144        self.assertEqual(1, len(m.get_transitions(source=self.States.RED, dest=self.States.YELLOW)))
145        self.assertEqual(9, len(m.get_transitions()))
146        m.add_transition('switch_to_yellow', self.States.RED, self.States.YELLOW)
147        self.assertEqual(4, len(m.get_transitions(source=self.States.RED)))
148        # we expect two return values. 'switch_to_yellow' and 'to_YELLOW'
149        self.assertEqual(2, len(m.get_transitions(source=self.States.RED, dest=self.States.YELLOW)))
150
151    def test_get_triggers(self):
152        m = self.machine_cls(states=self.States, initial=self.States.RED)
153        trigger_name = m.get_triggers(m.state.name)
154        trigger_enum = m.get_triggers(m.state)
155        self.assertEqual(trigger_enum, trigger_name)
156
157
158@skipIf(enum is None, "enum is not available")
159class TestNestedStateEnums(TestEnumsAsStates):
160
161    def setUp(self):
162        super(TestNestedStateEnums, self).setUp()
163        self.machine_cls = MachineFactory.get_predefined(nested=True)
164
165    def test_root_enums(self):
166        states = [self.States.RED, self.States.YELLOW,
167                  {'name': self.States.GREEN, 'children': ['tick', 'tock'], 'initial': 'tick'}]
168        m = self.machine_cls(states=states, initial=self.States.GREEN)
169        self.assertTrue(m.is_GREEN(allow_substates=True))
170        self.assertTrue(m.is_GREEN_tick())
171        m.to_RED()
172        self.assertTrue(m.state is self.States.RED)
173
174    def test_nested_enums(self):
175        states = ['A', self.States.GREEN,
176                  {'name': 'C', 'children': self.States, 'initial': self.States.GREEN}]
177        m1 = self.machine_cls(states=states, initial='C')
178        m2 = self.machine_cls(states=states, initial='A')
179        self.assertEqual(m1.state, self.States.GREEN)
180        self.assertTrue(m1.is_GREEN())  # even though it is actually C_GREEN
181        m2.to_GREEN()
182        self.assertTrue(m2.is_C_GREEN())  # even though it is actually just GREEN
183        self.assertEqual(m1.state, m2.state)
184        m1.to_A()
185        self.assertNotEqual(m1.state, m2.state)
186
187    def test_initial_enum(self):
188        m1 = self.machine_cls(states=self.States, initial=self.States.GREEN)
189        self.assertEqual(self.States.GREEN, m1.state)
190        self.assertEqual(m1.state.name, self.States.GREEN.name)
191
192    def test_duplicate_states(self):
193        with self.assertRaises(ValueError):
194            self.machine_cls(states=['A', 'A'])
195
196    def test_duplicate_states_from_enum_members(self):
197        class Foo(enum.Enum):
198            A = 1
199
200        with self.assertRaises(ValueError):
201            self.machine_cls(states=[Foo.A, Foo.A])
202
203    def test_add_enum_transition(self):
204
205        class Foo(enum.Enum):
206            A = 0
207            B = 1
208
209        class Bar(enum.Enum):
210            FOO = Foo
211            C = 2
212
213        m = self.machine_cls(states=Bar, initial=Bar.C, auto_transitions=False)
214        m.add_transition('go', Bar.C, Foo.A, conditions=lambda: False)
215        trans = m.events['go'].transitions['C']
216        self.assertEqual(1, len(trans))
217        self.assertEqual('FOO_A', trans[0].dest)
218        m.add_transition('go', Bar.C, 'FOO_B')
219        self.assertEqual(2, len(trans))
220        self.assertEqual('FOO_B', trans[1].dest)
221        m.go()
222        self.assertTrue(m.is_FOO_B())
223        m.add_transition('go', Foo.B, 'C')
224        trans = m.events['go'].transitions['FOO_B']
225        self.assertEqual(1, len(trans))
226        self.assertEqual('C', trans[0].dest)
227        m.go()
228        self.assertEqual(m.state, Bar.C)
229
230    def test_add_nested_enums_as_nested_state(self):
231        class Foo(enum.Enum):
232            A = 0
233            B = 1
234
235        class Bar(enum.Enum):
236            FOO = Foo
237            C = 2
238
239        m = self.machine_cls(states=Bar, initial=Bar.C)
240        self.assertEqual(sorted(m.states['FOO'].states.keys()), ['A', 'B'])
241        m.add_transition('go', 'FOO_A', 'C')
242        m.add_transition('go', 'C', 'FOO_B')
243        m.add_transition('foo', Bar.C, Bar.FOO)
244
245        m.to_FOO_A()
246        self.assertFalse(m.is_C())
247        self.assertTrue(m.is_FOO(allow_substates=True))
248        self.assertTrue(m.is_FOO_A())
249        self.assertTrue(m.is_FOO_A(allow_substates=True))
250        m.go()
251        self.assertEqual(Bar.C, m.state)
252        m.go()
253        self.assertEqual(Foo.B, m.state)
254        m.to_state(m, Bar.C.name)
255        self.assertEqual(Bar.C, m.state)
256        m.foo()
257        self.assertEqual(Bar.FOO, m.state)
258
259    def test_enum_model_conversion(self):
260        class Inner(enum.Enum):
261            I1 = 1
262            I2 = 2
263            I3 = 3
264            I4 = 0
265
266        class Middle(enum.Enum):
267            M1 = 10
268            M2 = 20
269            M3 = 30
270            M4 = Inner
271
272        class Outer(enum.Enum):
273            O1 = 100
274            O2 = 200
275            O3 = 300
276            O4 = Middle
277
278        m = self.machine_cls(states=Outer, initial=Outer.O1)
279
280    def test_enum_initial(self):
281        class Foo(enum.Enum):
282            A = 0
283            B = 1
284
285        class Bar(enum.Enum):
286            FOO = dict(children=Foo, initial=Foo.A)
287            C = 2
288
289        m = self.machine_cls(states=Bar, initial=Bar.FOO)
290        self.assertTrue(m.is_FOO_A())
291
292    def test_separator_naming_error(self):
293        class UnderscoreEnum(enum.Enum):
294            STATE_NAME = 0
295
296        # using _ in enum names in the default config should raise an error
297        with self.assertRaises(ValueError):
298            self.machine_cls(states=UnderscoreEnum)
299
300        # changing the separator should make it work
301        class DotNestedState(self.machine_cls.state_cls):
302            separator = '.'
303
304        # make custom machine use custom state with dot separator
305        class DotMachine(self.machine_cls):
306            state_cls = DotNestedState
307
308        m = DotMachine(states=UnderscoreEnum)
309
310    def test_get_nested_transitions(self):
311
312        class Errors(enum.Enum):
313            NONE = self.States
314            UNKNOWN = 2
315            POWER = 3
316        m = self.machine_cls(states=Errors, initial=Errors.NONE.value.RED, auto_transitions=False)
317        m.add_transition('error', Errors.NONE, Errors.UNKNOWN)
318        m.add_transition('outage', [Errors.NONE, Errors.UNKNOWN], Errors.POWER)
319        m.add_transition('reset', '*', self.States.RED)
320        m.add_transition('toggle', self.States.RED, self.States.GREEN)
321        m.add_transition('toggle', self.States.GREEN, self.States.YELLOW)
322        m.add_transition('toggle', self.States.YELLOW, self.States.RED)
323        self.assertEqual(5, len(m.get_transitions(dest=self.States.RED)))
324        self.assertEqual(1, len(m.get_transitions(source=self.States.RED, dest=self.States.RED, delegate=True)))
325        self.assertEqual(1, len(m.get_transitions(source=self.States.RED, dest=self.States.GREEN)))
326        self.assertEqual(1, len(m.get_transitions(dest=self.States.GREEN)))
327        self.assertEqual(3, len(m.get_transitions(trigger='toggle')))
328
329    def test_multiple_deeper(self):
330
331        class X(enum.Enum):
332            X1 = 1
333            X2 = 2
334
335        class B(enum.Enum):
336            B1 = dict(parallel=X)
337            B2 = 2
338
339        class A(enum.Enum):
340            A1 = dict(parallel=B)
341            A2 = 2
342
343        class Q(enum.Enum):
344            Q1 = 1
345            Q2 = dict(parallel=A)
346
347        class P(enum.Enum):
348            P1 = 1
349            P2 = dict(parallel=Q)
350
351        class States(enum.Enum):
352            S1 = 1
353            S2 = dict(parallel=P)
354
355        m = self.machine_cls(states=States, initial=States.S1)
356        self.assertEqual(m.state, States.S1)
357        m.to_S2()
358
359        ref_state = [P.P1, [Q.Q1, [[[X.X1, X.X2], B.B2], A.A2]]]
360        self.assertEqual(ref_state, m.state)
361
362
363@skipIf(enum is None or (pgv is None and gv is None), "enum and (py)graphviz are not available")
364class TestEnumWithGraph(TestEnumsAsStates):
365
366    def setUp(self):
367        super(TestEnumWithGraph, self).setUp()
368        self.machine_cls = MachineFactory.get_predefined(graph=True)
369
370    def test_get_graph(self):
371        m = self.machine_cls(states=self.States, initial=self.States.GREEN)
372        roi = m.get_graph(show_roi=False)
373        self.assertIsNotNone(roi)
374
375    def test_get_graph_show_roi(self):
376        m = self.machine_cls(states=self.States, initial=self.States.GREEN)
377        roi = m.get_graph(show_roi=True)
378        self.assertIsNotNone(roi)
379
380
381@skipIf(enum is None or (pgv is None and gv is None), "enum and (py)graphviz are not available")
382class TestNestedStateGraphEnums(TestNestedStateEnums):
383
384    def setUp(self):
385        super(TestNestedStateGraphEnums, self).setUp()
386        self.machine_cls = MachineFactory.get_predefined(nested=True, graph=True)
387