Better refactoring of message classes

lib/message/message.py

@@ -6,51 +6,29 @@ class MessageFormat():
     Defines interface and common methods for message formats.
     """
     required_params = []
-    nested_message_key = False
 
     def __init__(self, raw_message):
         self.raw_message = raw_message
 
-        self.message = self.get_message()
+        try:
+            self.json = json.loads(raw_message)
+        except:
+            self.json = False
 
-    @property
-    def missing(self):
-        if self.message:
-            return([x for x in self.required_params if x not in self.message])
-        else:
-            return self.required_params
+        self.message = self.get_message()
+        if not self.message:
+            raise TypeError("Message not of type '%s'" % (self.__class__.__name__))
 
     def get_message(self):
-        if isinstance(self.raw_message, str) or isinstance(self.raw_message, unicode):
-            try:
-                json_message = json.loads(self.raw_message)
-
-                if self.nested_message_key:
-                    if self.nested_message_key in json_message:
-                        return json.loads(json_message[self.nested_message_key])
-                    else:
-                        raise KeyError("%s not in nested message of type %s" % (self.nested_message_key, self.__class__.__name__))
-                else:
-                    return json_message
-            except ValueError:
-                return self.raw_message
-
-    def is_valid_format(self):
-        if "type" in self.message:
-            message_type = self.message['type']
-        elif "Type" in self.message:
-            message_type = self.message['Type']
+        if self.json:
+            return self.json
         else:
-            raise KeyError("Message missing 'Type' or 'type' key!")
+            return self.raw_message
 
-        if message_type == self._type:
-            return self.validate()
-        else:
-            return False
-
-    def validate(self):
-        if len(self.missing) > 0:
-            raise ValueError("Matched Type '%s', but missing required parameters: %s" % (self.__class__.__name__, self.missing))
+    def validate(self, message):
+        missing = [x for x in self.required_params if x not in message]
+        if len(missing) > 0:
+            raise ValueError("Matched Type '%s', but missing required parameters: %s" % (self.__class__.__name__, missing))
         else:
             return True
 
@@ -62,15 +40,32 @@ class RegistrationMessage(MessageFormat):
     _type = 'registration'
     required_params = ["method", "nagios_name", "chef_name"]
 
+    def get_message(self):
+        if self.json:
+            if 'type' in self.json and self.json['type'] == self._type:
+                if self.validate(self.json):
+                    return self.json
+
+        return False
+
 
 class AutoscalingMessage(MessageFormat):
     """
     Class to handle messages from Amazon's Autoscaling -> SNS -> SQS workflow.
     """
     _type = 'Notification'
-    nested_message_key = "Message"
     required_params = ["EC2InstanceId", "Event", "AutoScalingGroupName"]
 
+    def get_message(self):
+        if self.json:
+            if "Type" in self.json and self.json["Type"] == self._type:
+                if "Message" in self.json:
+                    message = json.loads(self.json["Message"])
+                    if self.validate(message):
+                        return message
+
+        return False
+
 
 class Message():
     """
@@ -84,10 +79,10 @@ def __init__(self, raw_message):
 
     def get_format(self, raw_message):
         for format in self.supported:
-            f = format(raw_message)
-
-            if f.is_valid_format():
-                return f
+            try:
+                return format(raw_message)
+            except TypeError:
+                continue
 
         raise ValueError("Message format unsupported.")