fix polymodel put and get

src/google/cloud/ndb/model.py

@@ -521,6 +521,14 @@ def _entity_from_ds_entity(ds_entity, model_class=None):
     """
     model_class = model_class or Model._lookup_model(ds_entity.kind)
     entity = model_class()
+
+    # Check if we are dealing with a PolyModel, and if so get correct subclass.
+    # We need to import here to avoid circular import.
+    from google.cloud.ndb import PolyModel
+
+    if isinstance(entity, PolyModel) and "class" in ds_entity:
+        entity = entity._class_map[tuple(ds_entity["class"])]()
+
     if ds_entity.key:
         entity._key = key_module.Key._from_ds_key(ds_entity.key)
 
@@ -655,8 +663,7 @@ def _entity_to_ds_entity(entity, set_key=True):
     if set_key:
         key = entity._key
         if key is None:
-            # use _class_name instead of _get_kind, to get PolyModel right
-            key = key_module.Key(entity._class_name(), None)
+            key = key_module.Key(entity._get_kind(), None)
         ds_entity = ds_entity_module.Entity(
             key._key, exclude_from_indexes=exclude_from_indexes
         )
@@ -1950,12 +1957,10 @@ def _validate_key(value, entity=None):
         raise exceptions.BadValueError("Expected Key, got {!r}".format(value))
 
     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
-        if value.kind() != entity._class_name():
+        if value.kind() != entity._get_kind():
             raise KindError(
                 "Expected Key kind to be {}; received "
-                "{}".format(entity._class_name(), value.kind())
+                "{}".format(entity._get_kind(), value.kind())
             )
 
     return value
@@ -4518,7 +4523,7 @@ class a different name when stored in Google Cloud Datastore than the
 
     @classmethod
     def _class_name(cls):
-        """A hook for polymodel to override.
+        """A hook for PolyModel to override.
 
         For regular models and expandos this is just an alias for
         _get_kind().  For PolyModel subclasses, it returns the class name

tests/system/test_query.py

@@ -528,6 +528,35 @@ def make_entities():
     assert not more
 
 
+@pytest.mark.usefixtures("client_context")
+def test_polymodel_query(ds_entity):
+    class Animal(ndb.PolyModel):
+        foo = ndb.IntegerProperty()
+
+    class Cat(Animal):
+        foo = ndb.IntegerProperty()
+
+    animal = Animal(foo=1)
+    animal.put()
+    cat = Cat(foo=2)
+    cat.put()
+
+    query = Animal.query()
+    results = eventually(query.fetch, _length_equals(2))
+
+    results = sorted(results, key=operator.attrgetter("foo"))
+    assert isinstance(results[0], Animal)
+    assert not isinstance(results[0], Cat)
+    assert isinstance(results[1], Animal)
+    assert isinstance(results[1], Cat)
+
+    query = Cat.query()
+    results = eventually(query.fetch, _length_equals(1))
+
+    assert isinstance(results[0], Animal)
+    assert isinstance(results[0], Cat)
+
+
 @pytest.mark.skip("Requires an index")
 @pytest.mark.usefixtures("client_context")
 def test_query_repeated_property(ds_entity):

tests/unit/test_model.py

@@ -1347,11 +1347,11 @@ class Mine(model.Model):
 
         value = model.Key(Mine, "yours")
         entity = unittest.mock.Mock(spec=Mine)
-        entity._class_name.return_value = "Mine"
+        entity._get_kind.return_value = "Mine"
 
         result = model._validate_key(value, entity=entity)
         assert result is value
-        entity._class_name.assert_called_once_with()
+        entity._get_kind.assert_called_once_with()
 
     @staticmethod
     @pytest.mark.usefixtures("in_context")
@@ -1361,13 +1361,13 @@ class Mine(model.Model):
 
         value = model.Key(Mine, "yours")
         entity = unittest.mock.Mock(spec=Mine)
-        entity._class_name.return_value = "NotMine"
+        entity._get_kind.return_value = "NotMine"
 
         with pytest.raises(model.KindError):
             model._validate_key(value, entity=entity)
 
         calls = [unittest.mock.call(), unittest.mock.call()]
-        entity._class_name.assert_has_calls(calls)
+        entity._get_kind.assert_has_calls(calls)
 
 
 class TestModelKey:

tests/unit/test_polymodel.py

@@ -108,8 +108,9 @@ class Animal(polymodel.PolyModel):
         class Cat(Animal):
             pass
 
-        key = datastore.Key("Cat", 123, project="testing")
+        key = datastore.Key("Animal", 123, project="testing")
         datastore_entity = datastore.Entity(key=key)
+        datastore_entity["class"] = ["Animal", "Cat"]
         protobuf = helpers.entity_to_protobuf(datastore_entity)
         entity = model._entity_from_protobuf(protobuf)
         assert isinstance(entity, Cat)