Fix psubscribe with handlers

fakeredis.py

@@ -1968,7 +1968,7 @@ def publish(self, channel, message):
                 del self._pubsubs[i]
                 continue
 
-            count += ps.put(channel, to_bytes(message), 'message')
+            count += ps.put(channel, message, 'message')
 
         return count
 
@@ -2209,9 +2209,9 @@ def __init__(self, decode_responses=False, *args, **kwargs):
         self.ignore_subscribe_messages = kwargs.get(
             'ignore_subscribe_messages', False)
 
-    def _normalize(self, channel):
-        channel = to_bytes(channel)
-        return _decode(channel) if self._decode_responses else channel
+    def _normalize(self, value):
+        value = to_bytes(value)
+        return _decode(value) if self._decode_responses else value
 
     def _normalize_keys(self, data):
         """
@@ -2232,6 +2232,7 @@ def put(self, channel, message, message_type):
             return self._send(message_type, None, channel, message)
 
         count = 0
+        message = self._normalize(message)
 
         # Send the message on the given channel
         if channel in self.channels:
@@ -2305,7 +2306,7 @@ def _subscribe(self, subscribed_dict, message_type, total_subscriptions,
                 new_channels[arg] = subscriber(arg, None)
 
         for channel, handler in iteritems(kwargs):
-            new_channels[channel] = handler
+            new_channels[channel] = subscriber(channel, handler)
 
         subscribed_dict.update(self._normalize_keys(new_channels))
         self.subscribed = True

test_fakeredis.py

@@ -2867,6 +2867,36 @@ def _listen(pubsub, q):
         self.assertEqual(msg4['data'], msg)
         self.assertIn(msg4['channel'], bpatterns)
 
+    @attr('slow')
+    def test_pubsub_listen_handler(self):
+        def _handler(message):
+            calls.append(message)
+
+        channel = 'ch1'
+        patterns = {'ch?': _handler}
+        calls = []
+
+        pubsub = self.redis.pubsub()
+        pubsub.subscribe(ch1=_handler)
+        pubsub.psubscribe(**patterns)
+        sleep(1)
+        msg1 = pubsub.get_message()
+        msg2 = pubsub.get_message()
+        self.assertEqual(msg1['type'], 'subscribe')
+        self.assertEqual(msg2['type'], 'psubscribe')
+        msg = 'hello world'
+        self.redis.publish(channel, msg)
+        sleep(1)
+        for i in range(2):
+            msg = pubsub.get_message()
+            self.assertIsNone(msg)   # get_message returns None when handler is used
+        pubsub.close()
+        calls.sort(key=lambda call: call['type'])
+        self.assertEqual(calls, [
+            {'pattern': None, 'channel': b'ch1', 'data': b'hello world', 'type': 'message'},
+            {'pattern': b'ch?', 'channel': b'ch1', 'data': b'hello world', 'type': 'pmessage'}
+        ])
+
     @attr('slow')
     def test_pubsub_ignore_sub_messages_listen(self):
         def _listen(pubsub, q):