Merge remote branch 'tom/hacks'

3 files changed, 169 insertions, 40 deletions

diff --git a/sleekxmpp/xmlstream/statemachine.py b/sleekxmpp/xmlstream/statemachine.py
index 0c841a9f..67b514a2 100644
--- a/sleekxmpp/xmlstream/statemachine.py
+++ b/sleekxmpp/xmlstream/statemachine.py
@@ -63,20 +63,22 @@ class StateMachine(object):
 		return self.transition_any( (from_state,), to_state, wait=wait, 
 		                            func=func, args=args, kwargs=kwargs )
 	
+	
 	def transition_any(self, from_states, to_state, wait=0.0, func=None, args=[], kwargs={} ):
 		'''
 		Transition from any of the given `from_states` to the given `to_state`.
 		'''
 
-		with self.lock:
-			for state in from_states:
-				if isinstance(state,tuple) or isinstance(state,list): 
-					raise ValueError( "State %s should be a string.  Did you mean to call 'StateMachine.transition_any()?" % str(state) )
-				if not state in self.__states: 
-					raise ValueError( "StateMachine does not contain from_state %s." % state )
-			if not to_state in self.__states: 
-				raise ValueError( "StateMachine does not contain to_state %s." % to_state )
+		if not (isinstance(from_states,tuple) or isinstance(from_states,list)): 
+				raise ValueError( "from_states should be a list or tuple" )
 
+		for state in from_states:
+			if not state in self.__states: 
+				raise ValueError( "StateMachine does not contain from_state %s." % state )
+		if not to_state in self.__states: 
+			raise ValueError( "StateMachine does not contain to_state %s." % to_state )
+
+		with self.lock:
 			start = time.time()
 			while not self.__current_state in from_states: 
 				# detect timeout:
@@ -95,30 +97,67 @@ class StateMachine(object):
 
 				logging.debug(' ==== TRANSITION %s -> %s', self.__current_state, to_state)
 				self.__current_state = to_state
-				self.lock.notifyAll()
+				self.lock.notify_all()
 				return return_val  # some 'true' value returned by func or True if func was None
 			else:
 				logging.error( "StateMachine bug!!  The lock should ensure this doesn't happen!" )
 				return False
 
 
+	def transition_ctx(self, from_state, to_state, wait=0.0):
+		'''
+		Use the state machine as a context manager.  The transition occurs on /exit/ from
+		the `with` context, so long as no exception is thrown.  For example:
+		
+		::
+
+			with state_machine.transition_ctx('one','two', wait=5) as locked:
+				if locked:
+					# the state machine is currently locked in state 'one', and will 
+					# transition to 'two' when the 'with' statement ends, so long as 
+					# no exception is thrown.
+					print 'Currently locked in state one: %s' % state_machine['one']
+
+				else:
+					# The 'wait' timed out, and no lock has been acquired
+					print 'Timed out before entering state "one"'
+
+			print 'Since no exception was thrown, we are now in state "two": %s' % state_machine['two']
+
+
+		The other main difference between this method and `transition()` is that the 
+		state machine is locked for the duration of the `with` statement.  Normally, 
+		after a `transition()` occurs, the state machine is immediately unlocked and 
+		available to another thread to call `transition()` again.
+		'''
+
+		if not from_state in self.__states: 
+			raise ValueError( "StateMachine does not contain from_state %s." % from_state )
+		if not to_state in self.__states: 
+			raise ValueError( "StateMachine does not contain to_state %s." % to_state )
+
+		return _StateCtx(self, from_state, to_state, wait)
+
+	
 	def ensure(self, state, wait=0.0):
 		'''
 		Ensure the state machine is currently in `state`, or wait until it enters `state`.
 		'''
 		return self.ensure_any( (state,), wait=wait )
 
+
 	def ensure_any(self, states, wait=0.0):
 		'''
 		Ensure we are currently in one of the given `states`
 		'''
-		with self.lock:
-			for state in states:
-				if isinstance(state,tuple) or isinstance(state,list): 
-					raise ValueError( "State %s should be a string.  Did you mean to call 'StateMachine.transition_any()?" % str(state) )
-				if not state in self.__states: 
-					raise ValueError( "StateMachine does not contain state %s." % state )
+		if not (isinstance(states,tuple) or isinstance(states,list)): 
+			raise ValueError('states arg should be a tuple or list')
 
+		for state in states:
+			if not state in self.__states: 
+				raise ValueError( "StateMachine does not contain state '%s'" % state )
+
+		with self.lock:
 			start = time.time()
 			while not self.__current_state in states: 
 				# detect timeout:
@@ -132,6 +171,11 @@ class StateMachine(object):
 		self.transition(self.__current_state, self._default_state)
 
 
+	def _set_state(self, state): #unsynchronized, only call internally after lock is acquired
+		self.__current_state = state
+		return state
+
+
 	def current_state(self):
 		'''
 		Return the current state name.
@@ -145,13 +189,46 @@ class StateMachine(object):
 		Use `StateMachine.ensure(state)` to wait until the machine enters a certain state.
 		'''
 		return self.__current_state == state
+
+	def __str__(self):
+		return "".join(( "StateMachine(", ','.join(self.__states), "): ", self.__current_state ))
+
 	
+
+class _StateCtx:
+
+	def __init__( self, state_machine, from_state, to_state, wait ):
+		self.state_machine = state_machine
+		self.from_state = from_state
+		self.to_state = to_state
+		self.wait = wait
+		self._timeout = False
+
 	def __enter__(self):
-		self.lock.acquire()
-		return self
+		self.state_machine.lock.acquire()
+		start = time.time()
+		while not self.state_machine[ self.from_state ]: 
+			# detect timeout:
+			if time.time() >= start + self.wait: 
+				logging.debug('StateMachine timeout while waiting for state: %s', self.from_state )
+				self._timeout = True # to indicate we should not transition
+				return False
+			self.state_machine.lock.wait(self.wait)
+
+		logging.debug('StateMachine entered context in state: %s', 
+				self.state_machine.current_state() )
+		return True
 
 	def __exit__(self, exc_type, exc_val, exc_tb):
-		self.lock.nofityAll()
-		self.lock.release()
+		if exc_val is not None:
+			logging.exception( "StateMachine exception in context, remaining in state: %s\n%s:%s", 
+				self.state_machine.current_state(), exc_type.__name__, exc_val )			
+		elif not self._timeout:
+			logging.debug(' ==== TRANSITION %s -> %s', 
+					self.state_machine.current_state(), self.to_state)
+			self.state_machine._set_state( self.to_state )
+
+		self.state_machine.lock.notify_all()
+		self.state_machine.lock.release()
 		return False # re-raise any exception
 
diff --git a/sleekxmpp/xmlstream/xmlstream.py b/sleekxmpp/xmlstream/xmlstream.py
index cc5e1ec8..6dbe7b30 100644
--- a/sleekxmpp/xmlstream/xmlstream.py
+++ b/sleekxmpp/xmlstream/xmlstream.py
@@ -101,19 +101,16 @@ class XMLStream(object):
 	
 	def connect(self, host='', port=0, use_ssl=None, use_tls=None):
 		"Establish a socket connection to the given XMPP server."
-		if not self.state.transition('disconnected','connecting'):
-			logging.warning("Can't connect now; Already in state %s", self.state.current_state())
+		
+		if not self.state.transition('disconnected','connected',
+				func=self.connectTCP, args=[host, port, use_ssl, use_tls] ):
+			
+			if self.state['connected']: logging.debug('Already connected')
+			else: logging.warning("Connection failed" )
 			return False
 
-		try:
-			return self.connectTCP(host, port, use_ssl, use_tls)
-		finally:
-			# attempt to ensure once a connection attempt starts, we leave either in the 
-			# 'connected' or 'disconnected' state.  Otherwise the connect method is not reentrant
-			if self.state['connecting']:
-				if not self.state.transition('connecting','disconnected'):
-					logging.error("Couldn't return to the 'disconnected' state after connection failure!")
-
+		logging.debug('Connection complete.')
+		return True
 
 		# TODO currently a caller can't distinguish between "connection failed" and
 		# "we're already trying to connect from another thread"
@@ -148,9 +145,6 @@ class XMLStream(object):
 				self.socket.connect(self.address)
 				self.filesocket = self.socket.makefile('rb', 0)
 				
-				if not self.state.transition('connecting','connected'):
-					logging.error( "State transition error!!!!  Shouldn't have happened" )
-				logging.debug('connect complete.')
 				return True
 
 			except socket.error as serr:
diff --git a/tests/test_statemachine.py b/tests/test_statemachine.py
index fec31098..e44b8e48 100644
--- a/tests/test_statemachine.py
+++ b/tests/test_statemachine.py
@@ -15,25 +15,23 @@ class testStateMachine(unittest.TestCase):
 	def testDefaults(self):
 		"Test ensure transitions occur correctly in a single thread"
 		s = sm.StateMachine(('one','two','three'))
-#		self.assertTrue(s.one)
 		self.assertTrue(s['one'])
-#		self.failIf(s.two)
 		self.failIf(s['two'])
 		try:
-			s.booga
+			s['booga']
 			self.fail('s.booga is an invalid state and should throw an exception!')
 		except: pass #expected exception
 
-	
+		# just make sure __str__ works, no reason to test its exact value:
+		print str(s)
+
+
 	def testTransitions(self):
 		"Test ensure transitions occur correctly in a single thread"
 		s = sm.StateMachine(('one','two','three'))
-#		self.assertTrue(s.one)
 
 		self.assertTrue( s.transition('one', 'two') )
-#		self.assertTrue( s.two )
 		self.assertTrue( s['two'] )
-#		self.failIf( s.one )
 		self.failIf( s['one'] )
 
 		self.assertTrue( s.transition('two', 'three') )
@@ -197,6 +195,66 @@ 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 result:
+			self.failIf( result )
+			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 result:
+				self.assertTrue( result )
+				self.assertTrue( s['two'] )
+				print 'thread 2 will transition on exit from the context manager...'
+			self.assertTrue( s['three'] )
+			print 'transitioned to %s' % s.current_state()
+
+		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)