python-statemachine/tests/examples/order_control_rich_model_machine.py at b39afce95dcf9deac9e6b10c4c5d97c8fb4ff191 · fgmacedo/python-statemachine · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
"""
Order control machine (rich model)
==================================

An StateMachine that demonstrates :ref:`Actions` being used on a rich model.

"""

from statemachine.exceptions import InvalidDefinition

from statemachine import State
from statemachine import StateMachine


class Order:
    def __init__(self):
        self.order_total = 0
        self.payments = []
        self.payment_received = False

    def payments_enough(self, amount):
        return sum(self.payments) + amount >= self.order_total

    def before_add_to_order(self, amount):
        self.order_total += amount
        return self.order_total

    def on_receive_payment(self, amount):
        self.payments.append(amount)
        return self.payments

    def after_receive_payment(self):
        self.payment_received = True

    def wait_for_payment(self):
        self.payment_received = False


class OrderControl(StateMachine):
    waiting_for_payment = State(initial=True, enter="wait_for_payment")
    processing = State()
    shipping = State()
    completed = State(final=True)

    add_to_order = waiting_for_payment.to(waiting_for_payment)
    receive_payment = waiting_for_payment.to(
        processing, cond="payments_enough"
    ) | waiting_for_payment.to(waiting_for_payment, unless="payments_enough")
    process_order = processing.to(shipping, cond="payment_received")
    ship_order = shipping.to(completed)


# %%
# Testing OrderControl
# --------------------
#
# Let's first try to create a statemachine instance, using the default dummy model that doesn't
# have the needed methods to complete the state machine. Since the required methods will not be
# found either in the state machine or in the model, an exception ``AttrNotFound`` will be raised.

try:
    control = OrderControl()
except InvalidDefinition as e:
    assert (  # noqa: PT017
        str(e)
        == (
            "Error on Waiting for payment when resolving callbacks: "
            "Did not found name 'wait_for_payment' from model or statemachine"
        )
    )


# %%
# Now initializing with a proper ``order`` instance.

order = Order()
control = OrderControl(order)

# %%
# Send events to add to order

assert control.send("add_to_order", 3) == 3
assert control.send("add_to_order", 7) == 10

# %%
# Receive a payment of $4...

control.send("receive_payment", 4)

# %%
# Since there's still $6 left to fulfill the payment, we cannot process the order.
try:
    control.send("process_order")
except StateMachine.TransitionNotAllowed as err:
    print(err)

# %%

control

# %%
# Now paying the left amount, we can proceed.

control.send("receive_payment", 6)

# %%

control

# %%

control.send("process_order")

# %%

control

# %%

control.send("ship_order")

# %%
# Just checking the final expected state

order.order_total

# %%

order.payments

# %%

control.completed.is_active

# %%

control


# %%
assert order.order_total == 10
assert order.payments == [4, 6]
assert control.completed.is_active