Handle exception when JSON received is not an object

Signed-off-by: Olivier Gayot <olivier.gayot@sigexec.com>

2 files changed, 23 insertions, 1 deletions

diff --git a/swiftstory/SwiftStory.py b/swiftstory/SwiftStory.py
index 770e347..dc55e2a 100644
--- a/swiftstory/SwiftStory.py
+++ b/swiftstory/SwiftStory.py
@@ -20,7 +20,11 @@ def message_received_handler(client, message):
     except json.JSONDecodeError:
         return error('badly formatted json')
 
-    op = json_msg.get('op')
+    try:
+        op = json_msg['op']
+    except (KeyError, TypeError):
+        op = None
+
     if op == 'join_game':
         try:
             game_name = json_msg['game_name']
diff --git a/tests/test_swiftstory.py b/tests/test_swiftstory.py
index 6a1f639..51dc837 100644
--- a/tests/test_swiftstory.py
+++ b/tests/test_swiftstory.py
@@ -11,6 +11,24 @@ class TestSwiftStory(unittest.TestCase):
                 SwiftStory.message_received_handler(client=None, message="{invalid_json}")
                 )
 
+    def test_receive_json_array(self):
+        self.assertEqual(
+                error("invalid command"),
+                SwiftStory.message_received_handler(client=None, message='[]')
+                )
+
+    def test_receive_json_number(self):
+        self.assertEqual(
+                error("invalid command"),
+                SwiftStory.message_received_handler(client=None, message='2.3')
+                )
+
+    def test_receive_json_null(self):
+        self.assertEqual(
+                error("invalid command"),
+                SwiftStory.message_received_handler(client=None, message='null')
+                )
+
     def test_receive_unknown_command(self):
         self.assertEqual(
                 error("invalid command"),