def test_as_dict(self):
"""Test use with a dict cache."""
self._db.basic.insert_many([
{'_id': 'A123', 'romeId': 'A123', 'name': 'Job Group 1'},
{'_id': 'A124', 'romeId': 'A124', 'name': 'Job Group 2'},
])
cache = {}
proto.cache_mongo_collection(self._db.basic.find, cache, job_pb2.JobGroup)
self.assertEqual(set(['A123', 'A124']), set(cache))
self.assertEqual('Job Group 2', cache['A124'].name)
def test_basic(self):
"""Test basic usage."""
self._db.basic.insert_many([
{'_id': 'A123', 'romeId': 'A123', 'name': 'Job Group 1'},
{'_id': 'A124', 'romeId': 'A124', 'name': 'Job Group 2'},
])
cache = []
proto.cache_mongo_collection(self._db.basic.find, cache, job_pb2.JobGroup)
self.assertEqual(['A123', 'A124'], [g.rome_id for g in cache])
self.assertEqual('A123', cache[0].rome_id)
self.assertEqual('Job Group 2', cache[1].name)
# Update the collection behind the scene.
self._db.basic.delete_one({'_id': 'A123'})
# The cache computed before the change should not have changed even if
# we call the function again.
proto.cache_mongo_collection(self._db.basic.find, cache, job_pb2.JobGroup)
self.assertEqual(['A123', 'A124'], [g.rome_id for g in cache])
# If the function is called with an empty cache it gets populated with
# the updated values.
cache2 = []
proto.cache_mongo_collection(self._db.basic.find, cache2, job_pb2.JobGroup)
self.assertEqual(['A124'], [g.rome_id for g in cache2])
def _chantiers():
"""Returns a list of known chantiers as protos."""
was_empty = not _CHANTIERS
all_chantiers = proto.cache_mongo_collection(_DB.chantiers.find,
_CHANTIERS,
chantier_pb2.Chantier)
if was_empty:
# Validate chantiers.
required_models = set(c.scoring_model
for c in all_chantiers.values()) | set(
scoring.GROUP_SCORING_MODELS.values())
existing_models = set(scoring.SCORING_MODELS) | set(
name
for name in required_models if scoring.get_scoring_model(name))
if required_models - existing_models:
logging.warning('Some scoring models will be random: %s',
required_models - existing_models)
if existing_models - required_models:
logging.warning('Some scoring models are unused: %s',
existing_models - required_models)
return all_chantiers
def _job_groups_info():
"""Returns a dict of info of known job groups as protos."""
return proto.cache_mongo_collection(_DB.job_group_info.find,
_JOB_GROUPS_INFO, job_pb2.JobGroup)
def templates(database):
"""Returns a list of known action templates as protos."""
return proto.cache_mongo_collection(database.action_templates.find,
_ACTION_TEMPLATES,
action_pb2.ActionTemplate)
def _sticky_action_steps(database):
"""Returns a dict of known sticky action steps keyed by ID."""
return proto.cache_mongo_collection(database.sticky_action_steps.find,
_STICKY_ACTION_STEPS,
action_pb2.StickyActionStep)
def _tip_templates(database):
"""Returns a list of known tip templates as protos."""
return proto.cache_mongo_collection(database.tip_templates.find,
_TIP_TEMPLATES,
action_pb2.ActionTemplate)
def _advice_modules(database):
return proto.cache_mongo_collection(database.advice_modules.find,
_ADVICE_MODULES,
advisor_pb2.AdviceModule)