diff options
| -rwxr-xr-x | redis/client.py | 6 | ||||
| -rw-r--r-- | tests/test_pubsub.py | 9 |
2 files changed, 15 insertions, 0 deletions
diff --git a/redis/client.py b/redis/client.py index 93da07c..0486022 100755 --- a/redis/client.py +++ b/redis/client.py @@ -3339,6 +3339,12 @@ class PubSub(object): ] self.reset() + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.reset() + def __del__(self): try: # if this object went out of scope prior to shutting down diff --git a/tests/test_pubsub.py b/tests/test_pubsub.py index 2644d60..31b60be 100644 --- a/tests/test_pubsub.py +++ b/tests/test_pubsub.py @@ -448,6 +448,15 @@ class TestPubSubAutoDecoding(object): new_data, pattern=self.pattern) + def test_context_manager(self, r): + with r.pubsub() as pubsub: + pubsub.subscribe('foo') + assert pubsub.connection is not None + + assert pubsub.connection is None + assert pubsub.channels == {} + assert pubsub.patterns == {} + class TestPubSubRedisDown(object): |