From 46dcde74549bdd73d5438f22ce066b09a2a66312 Mon Sep 17 00:00:00 2001
From: Blake Watters <blake@layer.com>
Date: Wed, 24 Jun 2015 15:25:52 -0700
Subject: [PATCH] Add strong enforcement that the `currentState` never
 transitions to `nil`

---
 Code/TKStateMachine.m      |  9 ++++++++-
 Specs/TKStateMachineSpec.m | 18 +++++++++++++++++-
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/Code/TKStateMachine.m b/Code/TKStateMachine.m
index 1edb75b..471ade0 100644
--- a/Code/TKStateMachine.m
+++ b/Code/TKStateMachine.m
@@ -104,6 +104,14 @@ - (void)setInitialState:(TKState *)initialState
     _initialState = initialState;
 }
 
+- (void)setCurrentState:(TKState *)currentState
+{
+    if (currentState == nil) {
+        [NSException raise:NSInvalidArgumentException format:@"Cannot assign currentState to `nil`: Expected a `TKState` object. (%@)", self];
+    }
+    _currentState = currentState;
+}
+
 - (NSSet *)states
 {
     return [NSSet setWithSet:self.mutableStates];
@@ -288,7 +296,6 @@ - (id)copyWithZone:(NSZone *)zone
 {
     TKStateMachine *copiedStateMachine = [[[self class] allocWithZone:zone] init];
     copiedStateMachine.active = NO;
-    copiedStateMachine.currentState = nil;
     copiedStateMachine.initialState = self.initialState;
     
     for (TKState *state in self.states) {
diff --git a/Specs/TKStateMachineSpec.m b/Specs/TKStateMachineSpec.m
index e1b0f3f..ffce220 100644
--- a/Specs/TKStateMachineSpec.m
+++ b/Specs/TKStateMachineSpec.m
@@ -57,6 +57,10 @@ - (void)startTryingToPickUpCollegeGirls {}
         [stateMachine.initialState shouldBeNil];
     });
     
+    it(@"has a nil current state", ^{
+        [stateMachine.currentState shouldBeNil];
+    });
+    
     it(@"has no events", ^{
         [[stateMachine.events should] haveCountOf:0];
     });
@@ -263,11 +267,23 @@ - (void)startTryingToPickUpCollegeGirls {}
         stateMachine.initialState = [stateMachine stateNamed:@"Dating"];
         [stateMachine activate];
     });
+    
+    context(@"when attempting to set the currentState to `nil`", ^{
+        it(@"should raise", ^{
+            __block NSException *exception;
+            @try {
+                [stateMachine setValue:nil forKey:@"currentState"];
+            } @catch (NSException *anException) {
+                exception = anException;
+            }
+            [[exception.name should] equal:NSInvalidArgumentException];
+            [[exception.reason should] startWithString:@"Cannot assign currentState to `nil`: Expected a `TKState` object."];
+        });
+    });
 
     it(@"invokes callbacks with a TKTransition describing the state change", ^{
         __block TKTransition *blockTransition;
         [singleState setWillEnterStateBlock:^(TKState *state, TKTransition *transition) {
-            NSLog(@"dsfdsfds");
             blockTransition = transition;
         }];
         NSError *error = nil;