Merge remote-tracking branch 'origin/features/object_discriminator'

Makefile

@@ -7,6 +7,7 @@ proto: groundstation/proto/gizmo_pb2.py \
 	groundstation/proto/object_list_pb2.py \
 	groundstation/proto/channel_list_pb2.py \
 	groundstation/proto/db_hash_pb2.py \
+	groundstation/objects/base_object_pb2.py \
 	groundstation/objects/root_object_pb2.py \
 	groundstation/objects/update_object_pb2.py
 

groundstation/objects/base_object.proto

@@ -0,0 +1,8 @@
+enum GizmoType {
+    UNSET = 0;
+    ROOT = 1;
+    UPDATE = 2;
+}
+message BaseObject {
+    required GizmoType type = 15;
+}

groundstation/objects/object_factory.py

@@ -1,8 +1,29 @@
+from groundstation.objects.base_object_pb2 import BaseObject, \
+    ROOT as TYPE_ROOT, \
+    UPDATE as TYPE_UPDATE, \
+    UNSET as TYPE_UNSET
 from groundstation.objects.root_object import RootObject
 from groundstation.objects.update_object import UpdateObject
 
+
+def type_of(protobuf):
+    base = BaseObject()
+    base.ParseFromString(protobuf)
+    return base.type
+
+
 def hydrate_object(protobuf):
-    obj = RootObject.from_object(protobuf)
-    if not obj.protocol:
-        obj = UpdateObject.from_object(protobuf)
-    return obj
+    # Test if it's strongly typed first
+    _type = type_of(protobuf)
+    if _type == TYPE_ROOT:
+        return RootObject.from_object(protobuf)
+    elif _type == TYPE_UPDATE:
+        return UpdateObject.from_object(protobuf)
+    elif _type == TYPE_UNSET:
+        # Use existing heuristic
+        obj = RootObject.from_object(protobuf)
+        if not obj.protocol:
+            obj = UpdateObject.from_object(protobuf)
+        return obj
+    else:
+        raise Exception("Unknown type; %s" % (str(_type)))

groundstation/objects/root_object.proto

@@ -1,5 +1,9 @@
+import "groundstation/objects/base_object.proto";
+
 message RootObject {
     required string id = 1;
     required string channel = 2;
     required string protocol = 3;
+
+    optional GizmoType type = 15 [default = ROOT];
 }

groundstation/objects/update_object.proto

@@ -1,4 +1,8 @@
+import "groundstation/objects/base_object.proto";
+
 message UpdateObject {
     repeated bytes parents = 1;
     required bytes data = 2;
+
+    optional GizmoType type = 15 [default = UPDATE];
 }

test/test_object_discrimination.py

@@ -0,0 +1,47 @@
+import unittest
+
+import groundstation.objects.object_factory as object_factory
+
+from groundstation.objects.base_object_pb2 import BaseObject, \
+    ROOT as TYPE_ROOT, \
+    UPDATE as TYPE_UPDATE, \
+    UNSET as TYPE_UNSET
+from groundstation.objects.root_object_pb2 import RootObject
+from groundstation.objects.update_object_pb2 import UpdateObject
+
+
+def new_root_object(weak):
+    root = RootObject()
+    root.id = "butts"
+    root.channel = "butts"
+    root.protocol = "butts"
+    if not weak:
+        root.type = TYPE_ROOT
+    return root
+
+
+def new_update_object(weak, parents=[]):
+    update = UpdateObject()
+    update.parents.extend(parents)
+    update.data = "butts"
+    if not weak:
+        update.type = TYPE_UPDATE
+    return update
+
+
+class TypeOfTestCase(unittest.TestCase):
+    def test_strongly_typed_root(self):
+        root_str = new_root_object(False).SerializeToString()
+        self.assertEqual(object_factory.type_of(root_str), TYPE_ROOT)
+
+    def test_weakly_typed_root(self):
+        root_str = new_root_object(True).SerializeToString()
+        self.assertEqual(object_factory.type_of(root_str), TYPE_UNSET)
+
+    def test_strongly_typed_update(self):
+        update_str = new_update_object(False).SerializeToString()
+        self.assertEqual(object_factory.type_of(update_str), TYPE_UPDATE)
+
+    def test_weakly_typed_update(self):
+        update_str = new_update_object(True).SerializeToString()
+        self.assertEqual(object_factory.type_of(update_str), TYPE_UNSET)