Order control machine (rich model)

A StateChart that demonstrates Actions being used on a rich model.

from statemachine.exceptions import InvalidDefinition

from statemachine import State
from statemachine import StateChart


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(StateChart):
    allow_event_without_transition = False
    enable_self_transition_entries = False

    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)

[4]

Since there’s still $6 left to fulfill the payment, we cannot process the order.

try:
    control.send("process_order")
except StateChart.TransitionNotAllowed as err:
    print(err)

Can't process_order when in Waiting for payment.

control

Now paying the left amount, we can proceed.

control.send("receive_payment", 6)

[4, 6]

control

control.send("process_order")

control

control.send("ship_order")

Just checking the final expected state

order.order_total

10

order.payments

[4, 6]

control.completed.is_active

True

control

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