From d4cd9de799726a8ed6f0d30e15c9bf80047d7de6 Mon Sep 17 00:00:00 2001
From: Miguel Grinberg <miguelgrinberg50@gmail.com>
Date: Sun, 19 Jul 2015 18:59:14 -0700
Subject: [PATCH] Return the rooms a client is in

---
 socketio/base_manager.py   |  8 ++++++++
 socketio/server.py         | 10 ++++++++++
 tests/test_base_manager.py |  8 ++++++++
 tests/test_server.py       | 12 ++++++++++++
 4 files changed, 38 insertions(+)

diff --git a/socketio/base_manager.py b/socketio/base_manager.py
index 64adad1..535137b 100755
--- a/socketio/base_manager.py
+++ b/socketio/base_manager.py
@@ -71,6 +71,14 @@ class BaseManager(object):
         except KeyError:
             pass
 
+    def get_rooms(self, sid, namespace):
+        """Return the rooms a client is in."""
+        r = []
+        for room_name, room in six.iteritems(self.rooms[namespace]):
+            if room_name is not None and sid in room and room[sid]:
+                r.append(room_name)
+        return r
+
     def emit(self, event, data, namespace, room=None, skip_sid=None,
              callback=None):
         """Emit a message to a single client, a room, or all the clients
diff --git a/socketio/server.py b/socketio/server.py
index 8908155..0e0969b 100755
--- a/socketio/server.py
+++ b/socketio/server.py
@@ -236,6 +236,16 @@ class Server(object):
         self.logger.info('room %s is closing [%s]', room, namespace)
         self.manager.close_room(namespace, room)
 
+    def rooms(self, sid, namespace=None):
+        """Return the rooms a client is in.
+
+        :param sid: Session ID of the client.
+        :param namespace: The Socket.IO namespace for the event. If this
+                          argument is omitted the default namespace is used.
+        """
+        namespace = namespace or '/'
+        return self.manager.get_rooms(sid, namespace)
+
     def handle_request(self, environ, start_response):
         """Handle an HTTP request from the client.
 
diff --git a/tests/test_base_manager.py b/tests/test_base_manager.py
index 88106f7..2ef5161 100755
--- a/tests/test_base_manager.py
+++ b/tests/test_base_manager.py
@@ -107,6 +107,14 @@ class TestBaseManager(unittest.TestCase):
     def test_close_invalid_room(self):
         self.bm.close_room('/foo', 'bar')
 
+    def test_rooms(self):
+        self.bm.connect('123', '/foo')
+        self.bm.enter_room('123', '/foo', 'bar')
+        r = self.bm.get_rooms('123', '/foo')
+        self.assertEqual(len(r), 2)
+        self.assertIn('123', r)
+        self.assertIn('bar', r)
+
     def test_emit_to_sid(self):
         self.bm.connect('123', '/foo')
         self.bm.connect('456', '/foo')
diff --git a/tests/test_server.py b/tests/test_server.py
index b4fcd5c..a21953b 100755
--- a/tests/test_server.py
+++ b/tests/test_server.py
@@ -94,6 +94,18 @@ class TestServer(unittest.TestCase):
         s.close_room('room')
         s.manager.close_room.assert_called_once_with('/', 'room')
 
+    def test_rooms(self, eio):
+        mgr = mock.MagicMock()
+        s = server.Server(client_manager_class=mgr)
+        s.rooms('123', namespace='/foo')
+        s.manager.get_rooms.assert_called_once_with('123', '/foo')
+
+    def test_rooms_default_namespace(self, eio):
+        mgr = mock.MagicMock()
+        s = server.Server(client_manager_class=mgr)
+        s.rooms('123')
+        s.manager.get_rooms.assert_called_once_with('123', '/')
+
     def test_handle_request(self, eio):
         s = server.Server()
         s.handle_request('environ', 'start_response')