Use correct class when deserializing a PolyModel entity. (#186)

Fixes #179.

Author: Chris Rossi

packages/google-cloud-ndb/google/cloud/ndb/model.py

@@ -527,7 +527,13 @@ def _entity_from_ds_entity(ds_entity, model_class=None):
     Returns:
         .Model: The deserialized entity.
     """
-    model_class = model_class or Model._lookup_model(ds_entity.kind)
+    class_key = ds_entity.get("class")
+    if class_key:
+        kind = class_key[-1]
+    else:
+        kind = ds_entity.kind
+
+    model_class = model_class or Model._lookup_model(kind)
     entity = model_class()
 
     # Check if we are dealing with a PolyModel, and if so get correct subclass.

packages/google-cloud-ndb/tests/system/test_crud.py

@@ -853,21 +853,24 @@ class SomeKind(ndb.Expando):
 @pytest.mark.usefixtures("client_context")
 def test_insert_polymodel(dispose_of):
     class Animal(ndb.PolyModel):
-        pass
+        one = ndb.StringProperty()
 
     class Feline(Animal):
-        pass
+        two = ndb.StringProperty()
 
     class Cat(Feline):
-        pass
+        three = ndb.StringProperty()
 
-    entity = Cat()
+    entity = Cat(one="hello", two="dad", three="i'm in jail")
     key = entity.put()
 
     retrieved = key.get()
 
     assert isinstance(retrieved, Animal)
     assert isinstance(retrieved, Cat)
+    assert retrieved.one == "hello"
+    assert retrieved.two == "dad"
+    assert retrieved.three == "i'm in jail"
 
     dispose_of(key._key)