Merge branch 'master' of github.com:exalearn/colmena

colmena/queue/redis.py

@@ -74,19 +74,17 @@ 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()
 
     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

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'