1 files changed, 57 insertions, 0 deletions
diff --git a/tests/test_statemachine.py b/tests/test_statemachine.py
index fec31098..4cfb50de 100644
--- a/tests/test_statemachine.py
+++ b/tests/test_statemachine.py
@@ -197,6 +197,63 @@ class testStateMachine(unittest.TestCase):
 		self.assertTrue(s['two'])
 
 
+	def testContextManager(self):
+
+		s = sm.StateMachine(('one','two','three'))
+
+		with s.transition_ctx('one','two'):
+			self.assertTrue( s['one'] )
+			self.failIf( s['two'] )
+
+		#successful transition b/c no exception was thrown
+		self.assertTrue( s['two'] )
+		self.failIf( s['one'] )
+
+		# failed transition because exception is thrown:
+		try:
+			with s.transition_ctx('two','three'):
+				raise Exception("boom!")
+			self.fail('exception expected')
+		except: pass
+
+		self.failIf( s.current_state() in ('one','three') )
+		self.assertTrue( s['two'] )
+
+	def testCtxManagerTransitionFailure(self):
+
+		s = sm.StateMachine(('one','two','three'))
+
+		with s.transition_ctx('two','three') as _s:
+			self.assertTrue( _s['one'] )
+			self.failIf( _s.current_state in ('two','three') )
+
+		self.assertTrue( _s['one'] )
+		
+		def r1():
+			print 'thread 1 started'
+			self.assertTrue( s.transition('one','two') )
+			print 'thread 1 transitioned'
+
+		def r2():
+			print 'thread 2 started'
+			self.failIf( s['two'] )
+			with s.transition_ctx('two','three', 10) as _s:
+				self.assertTrue( _s['two'] )
+				print 'thread 2 will transition on exit from the context manager...'
+			self.assertTrue( s['three'] )
+
+		t1 = threading.Thread(target=r1)
+		t2 = threading.Thread(target=r2)
+
+		t2.start() # this should block until r1 goes
+		time.sleep(1)
+		t1.start()
+
+		t1.join()
+		t2.join()
+
+		self.assertTrue( s['three'] )
+
 
 suite = unittest.TestLoader().loadTestsFromTestCase(testStateMachine)