From 5544c5ec74ce111f4312467ffa088f22b0913a61 Mon Sep 17 00:00:00 2001
From: Logan Ward <WardLT@users.noreply.github.com>
Date: Fri, 5 Apr 2024 10:28:54 -0400
Subject: [PATCH] Fix pickling for Redis (#134)

Was not calling the superclass
---
 colmena/queue/redis.py             |  8 +++-----
 colmena/queue/tests/test_queues.py | 10 ++++++++++
 2 files changed, 13 insertions(+), 5 deletions(-)

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'