From b4639c32de0250503322d7927910bf06d0e7b96b Mon Sep 17 00:00:00 2001 From: Chris Rossi Date: Mon, 9 Sep 2019 12:13:23 -0400 Subject: [PATCH] Use correct class when deserializing a PolyModel entity. Fixes #179. --- google/cloud/ndb/model.py | 8 +++++++- tests/system/test_crud.py | 11 +++++++---- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/google/cloud/ndb/model.py b/google/cloud/ndb/model.py index 72cfd499..f8cb96c7 100644 --- a/google/cloud/ndb/model.py +++ b/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. diff --git a/tests/system/test_crud.py b/tests/system/test_crud.py index b871303d..9deb282c 100644 --- a/tests/system/test_crud.py +++ b/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)