Adding Trillian handling for invalid types (#13)

chat_unifier/parsers/trillian/line_parser.py

@@ -3,6 +3,19 @@
 import urllib
 from xml.dom import minidom
 
+
+class Error(Exception):
+    pass
+
+
+class InvalidSessionType(Error):
+    pass
+
+
+class InvalidMessageType(Error):
+    pass
+
+
 SessionStartLine = collections.namedtuple(
     'SessionStartLine',
     field_names=['timestamp', 'medium', 'sender', 'recipient'])
@@ -38,10 +51,17 @@ def parse(line):
 
 
 def _parse_session_attributes(attributes):
-    if attributes[u'type'] == u'start':
+    try:
+        session_type = attributes[u'type']
+    except KeyError:
+        raise InvalidSessionType('Session element has no \'type\' attribute')
+
+    if session_type == u'start':
         return _parse_session_start(attributes)
-    else:
+    elif session_type == u'stop':
         return _parse_session_stop(attributes)
+    else:
+        raise InvalidSessionType('Unrecognized session type: %s' % session_type)
 
 
 def _parse_session_start(attributes):
@@ -61,10 +81,17 @@ def _parse_session_stop(attributes):
 
 
 def _parse_message_attributes(attributes):
-    if attributes[u'type'] == u'outgoing_privateMessage':
+    try:
+        message_type = attributes[u'type']
+    except KeyError:
+        raise InvalidMessageType('Message element has no \'type\' attribute')
+
+    if message_type == u'outgoing_privateMessage':
         return _parse_outgoing_message(attributes)
-    else:
+    elif message_type == u'incoming_privateMessage':
         return _parse_incoming_message(attributes)
+    else:
+        raise InvalidMessageType('Unrecognized message type: %s' % message_type)
 
 
 def _parse_outgoing_message(attributes):

tests/parsers/trillian/test_line_parser.py

@@ -28,7 +28,15 @@ def test_parses_valid_session_stop_line(self):
                 '<session type="stop" time="1132522555" medium="AIM" to="RemoteBuddy234" from="LocalUser789"/>'
             ))
 
-    def test_parsers_valid_outgoing_private_message_line(self):
+    def test_invalid_session_type_raises_exception(self):
+        with self.assertRaises(line_parser.InvalidSessionType):
+            line_parser.parse('<session type="dummy_invalid_type"/>')
+
+    def test_session_with_no_type_attribute_raises_exception(self):
+        with self.assertRaises(line_parser.InvalidSessionType):
+            line_parser.parse('<session />')
+
+    def test_parses_valid_outgoing_private_message_line(self):
         self.assertEqual(
             line_parser.OutgoingPrivateMessageLine(
                 timestamp=datetime.datetime(2005, 8, 25, 1, 23, 31),
@@ -41,7 +49,7 @@ def test_parsers_valid_outgoing_private_message_line(self):
                 '<message type="outgoing_privateMessage" time="1124933011" medium="AIM" to="RemoteBuddy888" from="LocalUser111" from_display="Me" text="do%20you%20want%20me%20to%20bring%20up%20the%20books%20tomorrow%3F"/>'
             ))
 
-    def test_parsers_valid_incoming_private_message_line(self):
+    def test_parses_valid_incoming_private_message_line(self):
         self.assertEqual(
             line_parser.IncomingPrivateMessageLine(
                 timestamp=datetime.datetime(2005, 8, 25, 1, 23, 45),
@@ -53,3 +61,11 @@ def test_parsers_valid_incoming_private_message_line(self):
             line_parser.parse(
                 '<message type="incoming_privateMessage" time="1124933025" medium="AIM" to="LocalUser222" from="RemoteBuddy555" from_display="Steve" text="hmm%2E%2E%2E%20no%20thanks"/>'
             ))
+
+    def test_invalid_message_type_raises_exception(self):
+        with self.assertRaises(line_parser.InvalidMessageType):
+            line_parser.parse('<message type="dummy_invalid_type" />')
+
+    def test_message_with_no_type_attribute_raises_exception(self):
+        with self.assertRaises(line_parser.InvalidMessageType):
+            line_parser.parse('<message />')