Generalize reading from socket

tests/integration/instances/virtualbox.py

@@ -50,40 +50,8 @@ class VirtualBoxInstance(Instance):
 
 	def boot(self):
 		self.machine.launch_vm_process(self.session, 'headless').wait_for_completion(-1)
-		self.console_output = self._read_console_output()
+		from ..tools import read_from_socket
-
+		self.console_output = read_from_socket(self.serial_port_path, 'INIT: Entering runlevel: 2', 20)
-	def _read_console_output(self):
-		import socket
-		import select
-		import errno
-		import sys
-		console = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-		console.connect(self.serial_port_path)
-		console.setblocking(0)
-
-		runlvl_check = 'INIT: Entering runlevel: 2'
-		output = ''
-		ptr = 0
-		continue_select = True
-		while continue_select:
-			read_ready, _, _ = select.select([console], [], [])
-			if console in read_ready:
-				while True:
-					try:
-						read = console.recv(1024)
-						output += read
-						sys.stdout.write(read)
-						if runlvl_check in output[ptr:]:
-							continue_select = False
-						else:
-							ptr = len(output) - len(runlvl_check)
-						break
-					except socket.error, e:
-						if e.errno != errno.EWOULDBLOCK:
-							raise Exception(e)
-						continue_select = False
-		console.close()
-		return output
 
 	def shutdown(self):
 		self.session.console.power_down().wait_for_completion(-1)

tests/integration/tools/__init__.py

@@ -33,6 +33,46 @@ def waituntil(predicate, timeout=5, interval=0.05):
 	return False
 
 
+def read_from_socket(socket_path, termination_string, timeout):
+		import socket
+		import select
+		import errno
+		console = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+		console.connect(socket_path)
+		console.setblocking(0)
+
+		output = ''
+		ptr = 0
+		continue_select = True
+		nooutput_for = 0
+		select_timeout = .5
+		while continue_select:
+			read_ready, _, _ = select.select([console], [], [], select_timeout)
+			if console in read_ready:
+				nooutput_for = 0
+				while True:
+					try:
+						output += console.recv(1024)
+						if termination_string in output[ptr:]:
+							continue_select = False
+						else:
+							ptr = len(output) - len(termination_string)
+						break
+					except socket.error, e:
+						if e.errno != errno.EWOULDBLOCK:
+							raise Exception(e)
+						continue_select = False
+			else:
+				nooutput_for += select_timeout
+			if nooutput_for > timeout:
+				from exceptions import SocketReadTimeout
+				msg = ('Reading from socket `{path}\' timed out after {seconds} seconds.'
+				       .format(path=socket_path, timeout=nooutput_for))
+				raise SocketReadTimeout(msg)
+		console.close()
+		return output
+
+
 def bootstrap(manifest, build_server):
 	from bootstrapvz.remote.build_servers import LocalBuildServer
 	if isinstance(build_server, LocalBuildServer):

tests/integration/tools/exceptions.py

@@ -0,0 +1,4 @@
+
+
+class SocketReadTimeout(Exception):
+	pass