def test_get_model_from_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(
FooModel(id='abc', app=feconf.OPPIA_PROJECT_ID, prop='123'),
job_utils.get_ndb_model_from_beam_entity(beam_entity))
def create_entities(count, id_or_name=False):
"""Creates a list of entities with random keys."""
if id_or_name:
ids_or_names = [uuid.uuid4().int & ((1 << 63) - 1) for _ in range(count)]
else:
ids_or_names = [str(uuid.uuid4()) for _ in range(count)]
keys = [types.Key(['EntityKind', x], project='project') for x in ids_or_names]
return [types.Entity(key) for key in keys]
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_get_beam_key_from_ndb_key(self) -> None:
beam_key = beam_datastore_types.Key(('FooModel', 'abc'),
project=feconf.OPPIA_PROJECT_ID,
namespace=self.namespace)
# We use private _from_ds_key here because it provides functionality
# for obtaining an NDB key from a Beam key, and writing it ourselves
# would be too complicated.
ndb_key = datastore_services.Key._from_ds_key(beam_key.to_client_key()) # pylint: disable=protected-access
self.assertEqual(job_utils.get_beam_key_from_ndb_key(ndb_key),
beam_key)
def get_beam_key_from_ndb_key(ndb_key):
"""Returns an Apache Beam key equivalent to the given NDB key.
Args:
ndb_key: datastore_services.Key. The NDB key.
Returns:
beam_datastore_types.Key. The Apache Beam key.
"""
return beam_datastore_types.Key(
ndb_key.flat(), project=ndb_key.app(), namespace=ndb_key.namespace())
def process(self, element, timestamp=None, *args, **kwargs):
from apache_beam.io.gcp.datastore.v1new import types
kind_id = timestamp.strftime('%Y%m%d-%H%M%S') + '-' + str(element.pop('indice'))
key = types.Key([self.kind, kind_id], project=self.project)
entity = types.Entity(key=key)
entity.set_properties({
'user_id': element['user_id'],
'item_id': element['item_id'],
'created': timestamp,
'rule': 'RULE_NAME',
})
yield 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 CreateEntity(elem):
"""Make a SignalQualitySampleStats entity from summarised stats."""
key = ds_types.Key(
[
'SignalQuality',
'{bot_group}/{bot}/{measurement}'.format(**elem[0].as_dict()),
'SignalQualitySampleStats',
'0', # version
],
project='chromeperf',
)
entity = ds_types.Entity(key)
properties = elem[1].as_dict()
properties['updated_time'] = datetime.datetime.now()
entity.set_properties(properties)
return entity
def test_get_model_from_beam_entity_with_time(self):
utcnow = datetime.datetime.utcnow()
beam_entity = beam_datastore_types.Entity(
beam_datastore_types.Key(
('CoreModel', 'abc'), project=feconf.OPPIA_PROJECT_ID,
namespace=self.namespace))
beam_entity.set_properties({
'prop': 3.14,
'created_on': utcnow.replace(tzinfo=datetime.timezone.utc),
'last_updated': None,
'deleted': False,
})
self.assertEqual(
CoreModel(
id='abc', project=feconf.OPPIA_PROJECT_ID, prop=3.14,
created_on=utcnow),
job_utils.get_ndb_model_from_beam_entity(beam_entity))
