def test_from_and_then_to_model(self):
model = FooModel(id='abc', project=feconf.OPPIA_PROJECT_ID, prop='123')
self.assertEqual(
model,
job_utils.get_ndb_model_from_beam_entity(
job_utils.get_beam_entity_from_ndb_model(model)))
def _read_from_datastore_handler(self, pickled_query):
"""XML-RPC handler for a ReadFromDatastore request.
Args:
pickled_query: str. The encoded Apache Beam query to respect.
Returns:
str. The list of all models encoded as a pickled list of Apache Beam
entities.
"""
with self._datastore_lock:
# TODO(#11475): This is very wasteful and inefficient, but good
# enough for unit tests and local development. When we move to Cloud
# NDB we'll have the genuine datastoreio module, and we'll be able
# to delete this code entirely, so it's OK to depend on this
# implementation until then.
ndb_models = list(datastore_services.query_everything().iter())
beam_query = pickle.loads(pickled_query)
job_utils.apply_query_to_models(beam_query, ndb_models)
return pickle.dumps([
job_utils.get_beam_entity_from_ndb_model(model)
for model in ndb_models
])
def test_get_beam_entity_from_model(self):
model = FooModel(id='abc', project=feconf.OPPIA_PROJECT_ID, prop='123')
beam_entity = job_utils.get_beam_entity_from_ndb_model(model)
self.assertEqual(beam_entity.key.path_elements, ('FooModel', 'abc'))
self.assertEqual(beam_entity.key.project, feconf.OPPIA_PROJECT_ID)
self.assertEqual(beam_entity.properties, {'prop': '123'})
def test_from_and_then_to_beam_entity(self):
beam_entity = beam_datastore_types.Entity(
beam_datastore_types.Key(('FooModel', 'abc'),
project=feconf.OPPIA_PROJECT_ID))
beam_entity.set_properties({'prop': '123'})
self.assertEqual(
beam_entity,
job_utils.get_beam_entity_from_ndb_model(
job_utils.get_ndb_model_from_beam_entity(beam_entity)))
def test_from_and_then_to_beam_entity(self):
beam_entity = beam_datastore_types.Entity(
beam_datastore_types.Key(
('CoreModel', 'abc'), project=feconf.OPPIA_PROJECT_ID))
beam_entity.set_properties({
'prop': 123,
'created_on': None,
'last_updated': None,
'deleted': False,
})
self.assertEqual(
beam_entity,
job_utils.get_beam_entity_from_ndb_model(
job_utils.get_ndb_model_from_beam_entity(beam_entity)))
def test_read_from_datastore(self):
query_everything = job_utils.get_beam_query_from_ndb_query(
datastore_services.query_everything())
model_list = [
self.create_model(base_models.BaseModel, id='a'),
self.create_model(base_models.BaseModel, id='b'),
self.create_model(base_models.BaseModel, id='c'),
]
self.put_multi(model_list)
with self.stub.context():
model_pcoll = (self.pipeline
| self.stub.ReadFromDatastore(query_everything))
self.assert_pcoll_equal(model_pcoll, [
job_utils.get_beam_entity_from_ndb_model(m) for m in model_list
])
