From 35f639f21f285897903d9805569af26c3333300a Mon Sep 17 00:00:00 2001
From: Zac Medico <zmedico@gentoo.org>
Date: Fri, 3 Sep 2010 14:23:58 -0700
Subject: Add support for a timeout argument to QueueScheduler.run() and use it
 in IpcDaemonTestCase to implement a 40 second timeout in test cases.

---
 pym/_emerge/PollScheduler.py                | 16 ++++++++-----
 pym/_emerge/QueueScheduler.py               | 35 ++++++++++++++++++++++++-----
 pym/_emerge/TaskScheduler.py                |  1 +
 pym/portage/tests/ebuild/test_ipc_daemon.py | 16 ++++++++++++-
 4 files changed, 57 insertions(+), 11 deletions(-)

diff --git a/pym/_emerge/PollScheduler.py b/pym/_emerge/PollScheduler.py
index f536de0fb..e8082fe21 100644
--- a/pym/_emerge/PollScheduler.py
+++ b/pym/_emerge/PollScheduler.py
@@ -205,7 +205,7 @@ class PollScheduler(object):
 		del self._poll_event_handlers[f]
 		del self._poll_event_handler_ids[reg_id]
 
-	def _schedule_wait(self, wait_ids, timeout=None):
+	def _schedule_wait(self, wait_ids=None, timeout=None):
 		"""
 		Schedule until wait_id is not longer registered
 		for poll() events.
@@ -220,16 +220,22 @@ class PollScheduler(object):
 			wait_ids = frozenset([wait_ids])
 
 		start_time = None
+		remaining_timeout = timeout
+		timed_out = False
 		if timeout is not None:
-			start_time = 1000 * time.time()
+			start_time = time.time()
 		try:
-			while wait_ids.intersection(handler_ids):
-				f, event = self._next_poll_event(timeout=timeout)
+			while (wait_ids is None and event_handlers) or \
+				(wait_ids is not None and wait_ids.intersection(handler_ids)):
+				f, event = self._next_poll_event(timeout=remaining_timeout)
 				handler, reg_id = event_handlers[f]
 				handler(f, event)
 				event_handled = True
 				if timeout is not None:
-					if 1000 * time.time() - start_time >= timeout:
+					elapsed_time = time.time() - start_time
+					remaining_timeout = (timeout - 1000 * elapsed_time)
+					if remaining_timeout <= 0:
+						timed_out = True
 						break
 		except StopIteration:
 			event_handled = True
diff --git a/pym/_emerge/QueueScheduler.py b/pym/_emerge/QueueScheduler.py
index 0e39d6ad3..1379ffc6f 100644
--- a/pym/_emerge/QueueScheduler.py
+++ b/pym/_emerge/QueueScheduler.py
@@ -1,6 +1,8 @@
 # Copyright 1999-2010 Gentoo Foundation
 # Distributed under the terms of the GNU General Public License v2
 
+import time
+
 from _emerge.PollScheduler import PollScheduler
 
 class QueueScheduler(PollScheduler):
@@ -28,13 +30,36 @@ class QueueScheduler(PollScheduler):
 	def remove(self, q):
 		self._queues.remove(q)
 
-	def run(self):
+	def clear(self):
+		for q in self._queues:
+			q.clear()
 
-		while self._schedule():
-			self._poll_loop()
+	def run(self, timeout=None):
 
-		while self._running_job_count():
-			self._poll_loop()
+		start_time = None
+		timed_out = False
+		remaining_timeout = timeout
+		if timeout is not None:
+			start_time = time.time()
+
+		while self._schedule():
+			self._schedule_wait(timeout=remaining_timeout)
+			if timeout is not None:
+				elapsed_time = time.time() - start_time
+				remaining_timeout = (timeout - 1000 * elapsed_time)
+				if remaining_timeout <= 0:
+					timed_out = True
+					break
+
+		if timeout is None or not timed_out:
+			while self._running_job_count():
+				self._schedule_wait(timeout=remaining_timeout)
+				if timeout is not None:
+					elapsed_time = time.time() - start_time
+					remaining_timeout = (timeout - 1000 * elapsed_time)
+					if remaining_timeout <= 0:
+						timed_out = True
+						break
 
 	def _schedule_tasks(self):
 		"""
diff --git a/pym/_emerge/TaskScheduler.py b/pym/_emerge/TaskScheduler.py
index f093e2848..83c0cbe96 100644
--- a/pym/_emerge/TaskScheduler.py
+++ b/pym/_emerge/TaskScheduler.py
@@ -17,6 +17,7 @@ class TaskScheduler(object):
 			max_jobs=max_jobs, max_load=max_load)
 		self.sched_iface = self._scheduler.sched_iface
 		self.run = self._scheduler.run
+		self.clear = self._scheduler.clear
 		self._scheduler.add(self._queue)
 
 	def add(self, task):
diff --git a/pym/portage/tests/ebuild/test_ipc_daemon.py b/pym/portage/tests/ebuild/test_ipc_daemon.py
index 0a9a9a929..dee61a2d2 100644
--- a/pym/portage/tests/ebuild/test_ipc_daemon.py
+++ b/pym/portage/tests/ebuild/test_ipc_daemon.py
@@ -3,6 +3,8 @@
 
 import shutil
 import tempfile
+import time
+import portage
 from portage import os
 from portage import _python_interpreter
 from portage.tests import TestCase
@@ -16,6 +18,8 @@ from _emerge.TaskScheduler import TaskScheduler
 
 class IpcDaemonTestCase(TestCase):
 
+	_SCHEDULE_TIMEOUT = 40000 # 40 seconds
+
 	def testIpcDaemon(self):
 		tmpdir = tempfile.mkdtemp()
 		try:
@@ -49,13 +53,23 @@ class IpcDaemonTestCase(TestCase):
 					args=[BASH_BINARY, "-c",
 					'"$PORTAGE_BIN_PATH"/ebuild-ipc exit %d' % exitcode],
 					env=env, scheduler=task_scheduler.sched_iface)
+
+				self.received_command = False
 				def exit_command_callback():
+					self.received_command = True
 					proc.cancel()
 					daemon.cancel()
+
 				exit_command.reply_hook = exit_command_callback
 				task_scheduler.add(daemon)
 				task_scheduler.add(proc)
-				task_scheduler.run()
+				start_time = time.time()
+				task_scheduler.run(timeout=self._SCHEDULE_TIMEOUT)
+				task_scheduler.clear()
+
+				self.assertEqual(self.received_command, True,
+					"command not received after %d seconds" % \
+					(time.time() - start_time,))
 				self.assertEqual(exit_command.exitcode, exitcode)
 		finally:
 			shutil.rmtree(tmpdir)
-- 
cgit v1.2.3-1-g7c22