diff --git a/colmena/queue/redis.py b/colmena/queue/redis.py index dbc9bd1..d7ca1b3 100644 --- a/colmena/queue/redis.py +++ b/colmena/queue/redis.py @@ -74,10 +74,9 @@ def __init__(self, self.connect() def __setstate__(self, state): - self.__dict__ = state + super().__setstate__(state) - # If you find the RedisClient placeholder, - # attempt to reconnect + # If you find the RedisClient placeholder, attempt to reconnect if self.redis_client == 'connected': self.redis_client = None self.connect() @@ -85,8 +84,7 @@ def __setstate__(self, state): def __getstate__(self): state = super().__getstate__() - # If connected, remove the unpicklable RedisClient and - # put a placeholder instead + # If connected, remove the unpicklable RedisClient and put a placeholder instead if self.is_connected: state['redis_client'] = 'connected' return state diff --git a/colmena/queue/tests/test_queues.py b/colmena/queue/tests/test_queues.py index 87e29d5..824ac80 100644 --- a/colmena/queue/tests/test_queues.py +++ b/colmena/queue/tests/test_queues.py @@ -1,5 +1,6 @@ """Tests across different queue implementations""" from multiprocessing import Pool +import pickle as pkl from pytest import fixture, raises, mark from redis import Redis, ConnectionError @@ -79,3 +80,12 @@ def test_kill_signal(queue): queue.send_kill_signal() with raises(KillSignalException): queue.get_task() + + +def test_pickle(queue): + """Ensure that we can still send and receive after pickling""" + queue: ColmenaQueues = pkl.loads(pkl.dumps(queue)) + + queue.send_inputs(1, method='test') + _, task = queue.get_task() + assert task.method == 'test'