Implement pre/post hook support

7 files changed, 113 insertions, 15 deletions

diff --git a/bin/send_msg b/bin/send_msg
index 19e88de..b0ea160 100755
--- a/bin/send_msg
+++ b/bin/send_msg
@@ -26,7 +26,7 @@ if __name__ == "__main__":
             'kwargs': {}
             }]
 
-        send_request(s, msg, guarantee=True, reply_requested=True, timeout=1)
+        send_request(s, msg, guarantee=True, reply_requested=True, timeout=10)
         print zmq.POLLOUT
         events = dict(poller.poll(500))
         print events
diff --git a/eventmq/client/jobs.py b/eventmq/client/jobs.py
index f82210d..9ec9a07 100644
--- a/eventmq/client/jobs.py
+++ b/eventmq/client/jobs.py
@@ -101,7 +101,8 @@ class Job(object):
         return f
 
 
-def job(func, broker_addr=None, queue=None, async=True, *args, **kwargs):
+def job(func, broker_addr=None, queue=None, async=True, *args,
+        **kwargs):
     """
     Functional decorator helper for creating a deferred eventmq job. See
     :class:`Job` for more information.
diff --git a/eventmq/client/messages.py b/eventmq/client/messages.py
index 617720a..9e49ea8 100644
--- a/eventmq/client/messages.py
+++ b/eventmq/client/messages.py
@@ -108,7 +108,7 @@ def schedule(socket, func, interval_secs=None, args=(), kwargs=None,
 
 
 def defer_job(
-        socket, func, wrapper=None, args=(), kwargs=None, class_args=(),
+        socket, func, args=(), kwargs=None, class_args=(),
         class_kwargs=None, reply_requested=False, guarantee=False,
         retry_count=0, timeout=0, debounce_secs=False,
         queue=conf.DEFAULT_QUEUE_NAME):
@@ -126,8 +126,6 @@ def defer_job(
         socket (socket): eventmq socket to use for sending the message
         func (callable or str): the callable (or string path to callable) to be
             deferred to a worker
-        wrapper (callable): optional wrapper for the call to func to be
-             wrapped with
         args (list): list of *args for the callable
         kwargs (dict): dict of **kwargs for the callable
         class_args (list): list of *args to pass to the the class when
@@ -177,15 +175,6 @@ def defer_job(
         logger.error('Encountered non-callable func: {}'.format(func))
         return
 
-    if wrapper and callable(wrapper):
-        # Prepend the original path and callable name to args
-        args = (path, callable_name, args)
-        callable_name = name_from_callable(wrapper)
-        path, callable_name = split_callable_name(callable_name)
-    elif wrapper:
-        logger.error('Encountered non-callable wrapper: {}'.format(wrapper))
-        return
-
     if not callable_name or not path:
         logger.error('Encountered invalid callable, will not proceed.')
         return
diff --git a/eventmq/conf.py b/eventmq/conf.py
index cec0249..4ad17f7 100644
--- a/eventmq/conf.py
+++ b/eventmq/conf.py
@@ -88,4 +88,8 @@ RQ_PASSWORD = ''
 
 MAX_JOB_COUNT = 1024
 
+# Path/Callable to run on start of a worker process
+SETUP_PATH = ''
+SETUP_CALLABLE = ''
+
 # }}}
diff --git a/eventmq/tests/test_jobmanager.py b/eventmq/tests/test_jobmanager.py
index db82843..f7d41c7 100644
--- a/eventmq/tests/test_jobmanager.py
+++ b/eventmq/tests/test_jobmanager.py
@@ -170,3 +170,8 @@ def start_jm(jm, addr):
 
 def pretend_job(t):
     time.sleep(t)
+
+
+def test_setup():
+    import time
+    assert time
diff --git a/eventmq/tests/test_worker.py b/eventmq/tests/test_worker.py
new file mode 100644
index 0000000..2f2c67e
--- /dev/null
+++ b/eventmq/tests/test_worker.py
@@ -0,0 +1,65 @@
+# This file is part of eventmq.
+#
+# eventmq is free software: you can redistribute it and/or modify it under the
+# terms of the GNU Lesser General Public License as published by the Free
+# Software Foundation, either version 2.1 of the License, or (at your option)
+# any later version.
+#
+# eventmq is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with eventmq.  If not, see <http://www.gnu.org/licenses/>.
+
+from multiprocessing import Pool
+import time
+
+from nose import with_setup
+
+from .. import worker
+
+ADDR = 'inproc://pour_the_rice_in_the_thing'
+
+
+def setup_func():
+    global pool
+    global out
+    pool = Pool()
+    out = pool.map(job, range(1))
+
+
+@with_setup(setup_func)
+def test_run_with_timeout():
+    payload = {
+        'path': 'eventmq.tests.test_worker',
+        'callable': 'job',
+        'args': [2]
+    }
+
+    msgid = worker._run(payload)
+
+    assert msgid
+
+
+@with_setup(setup_func)
+def test_run_setup():
+    setup_callable = 'pre_hook'
+    setup_path = 'eventmq.tests.test_worker'
+
+    worker.run_setup(setup_path, setup_callable)
+
+
+def job(sleep_time=0):
+    time.sleep(sleep_time)
+
+    return True
+
+
+def pre_hook():
+    return 1
+
+
+def post_hook():
+    return 1
diff --git a/eventmq/worker.py b/eventmq/worker.py
index 8d25c63..1e8e5a4 100644
--- a/eventmq/worker.py
+++ b/eventmq/worker.py
@@ -49,11 +49,12 @@ class MultiprocessWorker(Process):
     Defines a worker that spans the job in a multiprocessing task
     """
 
-    def __init__(self, input_queue, output_queue):
+    def __init__(self, input_queue, output_queue, run_setup=True):
         super(MultiprocessWorker, self).__init__()
         self.input_queue = input_queue
         self.output_queue = output_queue
         self.job_count = 0
+        self.run_setup = run_setup
 
     def run(self):
         """
@@ -61,6 +62,15 @@ class MultiprocessWorker(Process):
 
         This is designed to run in a seperate process.
         """
+
+        if self.run_setup:
+            self.run_setup = False
+            if conf.SETUP_CALLABLE and conf.SETUP_PATH:
+                try:
+                    run_setup(conf.SETUP_PATH, conf.SETUP_CALLABLE)
+                except Exception as e:
+                    logger.warning('Unable to complete setup: ' + str(e))
+
         import zmq
         zmq.Context.instance().term()
 
@@ -104,6 +114,9 @@ class MultiprocessWorker(Process):
 
 
 def _run(payload):
+    """
+    Takes care of actually executing the code given a message payload
+    """
     if ":" in payload["path"]:
         _pkgsplit = payload["path"].split(':')
         s_package = _pkgsplit[0]
@@ -148,5 +161,26 @@ def _run(payload):
     except Exception as e:
         logger.exception(e)
         return str(e)
+
     # Signal that we're done with this job
     return 'DONE'
+
+
+def run_setup(setup_path, setup_callable):
+    logger.debug("Running setup for worker id {}".format(os.getpid()))
+    if ":" in setup_path:
+        _pkgsplit = setup_path.split(':')
+        s_setup_package = _pkgsplit[0]
+    else:
+        s_setup_package = setup_path
+
+    if setup_callable and s_setup_package:
+        setup_package = import_module(s_setup_package)
+
+        setup_callable_ = getattr(setup_package, setup_callable)
+
+        try:
+            setup_callable_()
+        except Exception as e:
+            logger.exception(e)
+            return str(e)