diff --git a/mrq/queue_raw.py b/mrq/queue_raw.py index ed681404..44fbc71b 100644 --- a/mrq/queue_raw.py +++ b/mrq/queue_raw.py @@ -195,7 +195,7 @@ def dequeue_jobs(self, max_jobs=1, job_class=None, worker=None): if worker: worker.status = "spawn" - job_data = [job_factory(p) for p in params] + job_data = [job_factory(*p) for p in self.get_job_factory_params(job_factory, params)] for j in job_data: j["status"] = "started" j["queue"] = retry_queue @@ -206,6 +206,22 @@ def dequeue_jobs(self, max_jobs=1, job_class=None, worker=None): for job in job_class.insert(job_data, statuses_no_storage=statuses_no_storage): yield job + def get_job_factory_params(self, job_factory, params): + args_count = job_factory.__code__.co_argcount + + if args_count == 0: + raise Exception("Job factory must take at least 1 argument!") + + if args_count == 1: + # just expecting raw params + return ([p] for p in params) + + if args_count == 2: + # expecting raw queue name as second argument + return ([p, self.id] for p in params) + + raise Exception("Job factory expecting too many arguments!") + def get_sorted_graph( self, start=0, diff --git a/tests/fixtures/config-raw1.py b/tests/fixtures/config-raw1.py index 79039f82..40164e1e 100644 --- a/tests/fixtures/config-raw1.py +++ b/tests/fixtures/config-raw1.py @@ -147,5 +147,14 @@ "sleep": float(rawparam.split(" ")[1]) } } + }, + "test_factory_params_raw": { + "job_factory": lambda rawparam, queue: { + "path": "tests.tasks.general.MongoInsert", + "params": { + "queue": queue, + "sleep": float(rawparam) + } + } } } diff --git a/tests/test_raw.py b/tests/test_raw.py index f69f88ed..32dab340 100644 --- a/tests/test_raw.py +++ b/tests/test_raw.py @@ -360,3 +360,26 @@ def test_raw_no_storage(worker): time.sleep(2) assert jobs_collection.count({"status": "retry"}) == 1 + + +def test_raw_factory_params(worker): + + p_queue = "test_factory_params_raw" + + worker.start( + flags="--greenlets 10 --config tests/fixtures/config-raw1.py", queues=p_queue) + + jobs_collection = worker.mongodb_jobs.mrq_jobs + tests_collection = worker.mongodb_jobs.tests_inserts + + assert jobs_collection.count() == 0 + assert tests_collection.count() == 0 + assert Queue(p_queue).size() == 0 + + worker.send_raw_tasks(p_queue, [0], block=True) + + assert Queue(p_queue).size() == 0 + assert jobs_collection.count() == 1 + assert jobs_collection.count({"status": "success"}) == 1 + assert tests_collection.count() == 1 + assert tests_collection.count({"params.queue": p_queue, "params.sleep": 0.0}) == 1