def test_get_identifier_mappings(self): """Test _get_identifier_mappings function.""" count = 3 with factories.single_commit(): regulation = factories.RegulationFactory() requirements = [factories.RequirementFactory() for _ in range(count)] controls = [factories.ControlFactory() for _ in range(count)] expected_id_map = { "Requirement": {o.id: o.slug for o in requirements}, "Control": {o.id: o.slug for o in controls}, } relationships = [] for i in range(count): relationships.append(factories.RelationshipFactory( source=regulation if i % 2 == 0 else requirements[i], destination=regulation if i % 2 == 1 else requirements[i], )) relationships.append(factories.RelationshipFactory( source=regulation if i % 2 == 0 else controls[i], destination=regulation if i % 2 == 1 else controls[i], )) block = base_block.ExportBlockConverter( mock.MagicMock(), object_class=models.Regulation, fields="all", object_ids=[regulation.id], class_name=models.Regulation.__name__, ) id_map = block._get_identifier_mappings(relationships) self.assertEqual(expected_id_map, id_map)
def initialize_block_converters(self): """Generate block converters. Generate block converters from a list of tuples with an object name and ids and store it to an instance variable. """ object_map = {o.__name__: o for o in self.exportable.values()} exportable_queries = self._get_exportable_queries() for object_data in exportable_queries: class_name = object_data["object_name"] object_class = object_map[class_name] object_ids = object_data.get("ids", []) fields = object_data.get("fields", "all") if class_name == "Snapshot": self.block_converters.append( SnapshotBlockConverter(self, object_ids, fields=fields) ) else: block_converter = base_block.ExportBlockConverter( self, object_class=object_class, fields=fields, object_ids=object_ids, class_name=class_name, ) self.block_converters.append(block_converter)
def test_create_mapping_cache(self, count): """Test creation of mapping cache for export.""" with factories.single_commit(): regulations = [factories.RegulationFactory() for _ in range(count)] requirements = [ factories.RequirementFactory() for _ in range(count) ] controls = [factories.ControlFactory() for _ in range(count)] expected_cache = defaultdict(lambda: defaultdict(list)) for i in range(count): for j in range(i): factories.RelationshipFactory( source=regulations[j] if i % 2 == 0 else requirements[i], destination=regulations[j] if i % 2 == 1 else requirements[i], ) with mock.patch( 'ggrc.models.relationship.is_external_app_user', return_value=True): factories.RelationshipFactory( source=regulations[j] if i % 2 == 0 else controls[i], destination=regulations[j] if i % 2 == 1 else controls[i], is_external=True, ) expected_cache[regulations[j].id]["Control"].append( controls[i].slug) expected_cache[regulations[j].id]["Requirement"].append( requirements[i].slug) block = base_block.ExportBlockConverter( mock.MagicMock(), object_class=models.Regulation, fields="all", object_ids=[r.id for r in regulations], class_name=models.Regulation.__name__, ) with QueryCounter() as counter: cache = block._create_mapping_cache() self.assertEqual( self.dd_to_dict(cache), self.dd_to_dict(expected_cache), ) self.assertLess(counter.get, self.QUERY_LIMIT)