1 files changed, 2 insertions, 10 deletions
diff --git a/examples/statemachine/trafficLightDemo.py b/examples/statemachine/trafficLightDemo.py
index 4d70541..8358750 100644
--- a/examples/statemachine/trafficLightDemo.py
+++ b/examples/statemachine/trafficLightDemo.py
@@ -8,21 +8,13 @@ import statemachine
 import trafficlightstate
 
 
-class TrafficLight:
+class TrafficLight(trafficlightstate.TrafficLightStateMixin):
     def __init__(self):
-        # start light in Red state
-        self._state = trafficlightstate.Red()
+        self.initialize_state(trafficlightstate.Red())
 
     def change(self):
         self._state = self._state.next_state()
 
-    # get light behavior/properties from current state
-    def __getattr__(self, attrname):
-        return getattr(self._state, attrname)
-
-    def __str__(self):
-        return "{0}: {1}".format(self.__class__.__name__, self._state)
-
 
 light = TrafficLight()
 for i in range(10):