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)