From 626499cd4596c13a41a5ecbe3e350b243c5fe6be Mon Sep 17 00:00:00 2001 From: Patrick Decat Date: Tue, 12 Apr 2016 11:58:51 +0200 Subject: [PATCH] Avoid KeyError when no room exists --- socketio/base_manager.py | 9 ++++++--- tests/test_base_manager.py | 4 ++++ 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/socketio/base_manager.py b/socketio/base_manager.py index d41ed609..cdbaa801 100644 --- a/socketio/base_manager.py +++ b/socketio/base_manager.py @@ -83,9 +83,12 @@ def close_room(self, room, namespace): 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) + try: + 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) + except KeyError: + pass return r def emit(self, event, data, namespace, room=None, skip_sid=None, diff --git a/tests/test_base_manager.py b/tests/test_base_manager.py index 36c10862..b916ddc9 100644 --- a/tests/test_base_manager.py +++ b/tests/test_base_manager.py @@ -131,6 +131,10 @@ def test_leave_invalid_room(self): self.bm.leave_room('123', '/foo', 'baz') self.bm.leave_room('123', '/bar', 'baz') + def test_no_room(self): + rooms = self.bm.get_rooms('123', '/foo') + self.assertEqual([], rooms) + def test_close_room(self): self.bm.connect('123', '/foo') self.bm.connect('456', '/foo')