add system test for PolyModel (#133)

src/google/cloud/ndb/__init__.py

@@ -186,6 +186,7 @@
 from google.cloud.ndb.model import User
 from google.cloud.ndb.model import UserNotFoundError
 from google.cloud.ndb.model import UserProperty
+from google.cloud.ndb.polymodel import PolyModel
 from google.cloud.ndb.query import ConjunctionNode
 from google.cloud.ndb.query import AND
 from google.cloud.ndb.query import DisjunctionNode

src/google/cloud/ndb/model.py

@@ -642,7 +642,8 @@ def _entity_to_ds_entity(entity, set_key=True):
     if set_key:
         key = entity._key
         if key is None:
-            key = key_module.Key(entity._get_kind(), None)
+            # use _class_name instead of _get_kind, to get PolyModel right
+            key = key_module.Key(entity._class_name(), None)
         ds_entity = entity_module.Entity(
             key._key, exclude_from_indexes=exclude_from_indexes
         )
@@ -1937,7 +1938,7 @@ def _validate_key(value, entity=None):
 
     if entity and type(entity) not in (Model, Expando):
         # Need to use _class_name instead of _get_kind, to be able to
-        # return the correct kind if this is a polymodel
+        # return the correct kind if this is a PolyModel
         if value.kind() != entity._class_name():
             raise KindError(
                 "Expected Key kind to be {}; received "

tests/system/test_crud.py

@@ -544,3 +544,25 @@ class SomeKind(ndb.Expando):
     assert retrieved.expando_prop == "exp-value"
 
     dispose_of(key._key)
+
+
+@pytest.mark.usefixtures("client_context")
+def test_insert_polymodel(dispose_of):
+    class Animal(ndb.PolyModel):
+        pass
+
+    class Feline(Animal):
+        pass
+
+    class Cat(Feline):
+        pass
+
+    entity = Cat()
+    key = entity.put()
+
+    retrieved = key.get()
+
+    assert isinstance(retrieved, Animal)
+    assert isinstance(retrieved, Cat)
+
+    dispose_of(key._key)