def test_constructor(self):
reader = Reader(self.source, self.app, self.dataset())
self.assertEqual(reader.app, self.app)
self.assertEqual(reader.source, self.source)
self.assertEqual(reader.datastore, None)
self.assertEqual(ImportMode.FULL, reader.mode)
reader = Reader(self.source, self.app, self.dataset(), 'other mode')
self.assertEqual('other mode', reader.mode)
def test_protect_row(self):
reader = Reader(self.source, self.app, self.dataset())
reader.secure_attributes = ['attrB']
row = {'attrA': 'valA', 'attrB': 'valB'}
self.assertEqual({
'attrA': 'valA',
'attrB': 'read_protected(valB)',
}, reader._protect_row(row))
def test_set_secure_attributes(self):
reader = Reader(self.source, self.app, self.dataset())
mapping = {
'a': {
'source_mapping': 'any secure string'
},
'b': {
'source_mapping': {
'bronwaarde': 'any secure bronwaarde'
}
},
'c': {
'source_mapping': 'any string'
},
'd': {
'source_mapping': {
'bronwaarde': 'any bronwaarde'
}
},
}
attributes = {
'a': {
'type': 'GOB.SecureString'
},
'b': {
'type': 'GOB.JSON',
'secure': {
'bronwaarde': {
'type': 'GOB.SecureString',
'level': 5
}
}
},
'c': {
'type': 'GOB.String'
},
'd': {
'type': 'GOB.JSON',
'secure': {
'bronwaarde': {
'type': 'GOB.String',
'level': 5
}
}
}
}
reader.set_secure_attributes(mapping, attributes)
self.assertEqual(reader.secure_attributes,
['any secure string', 'any secure bronwaarde'])
def test_connect(self, mock_datastore_factory, mock_datastore_config):
reader = Reader(self.source, self.app, self.dataset())
# 1. Should use application config to connect
reader.source = {
'application_config': 'the application config',
'application': 'the application',
'read_config': {
'read': 'config'
},
}
reader.mode = 'the mode'
reader.connect()
self.assertEqual(mock_datastore_factory.get_datastore.return_value,
reader.datastore)
reader.datastore.connect.assert_called_once()
mock_datastore_factory.get_datastore.assert_called_with(
'the application config', {
'read': 'config',
'mode': 'the mode'
})
mock_datastore_config.assert_not_called()
# 2. Application defined, no application_config defined. Should get application config
mock_datastore_factory.reset_mock()
reader.source = {'application': 'the application'}
reader.datastore = None
reader.connect()
self.assertEqual(mock_datastore_factory.get_datastore.return_value,
reader.datastore)
reader.datastore.connect.assert_called_once()
mock_datastore_factory.get_datastore.assert_called_with(
mock_datastore_config.return_value, {'mode': 'the mode'})
mock_datastore_config.assert_called_with('the application')
def import_rows(self, write, progress):
self.logger.info(f"Connect to {self.source_app}")
reader = Reader(self.source, self.source_app, self.dataset, self.mode)
reader.connect()
self.logger.info(f"Start import from {self.source_app}")
self.n_rows = 0
for row in reader.read():
progress.tick()
self.row = row
self.n_rows += 1
self.injector.inject(row)
self.enricher.enrich(row)
self.merger.merge(row, write)
entity = self.converter.convert(row)
# validator and entity_validator build up sets of primary keys from the dataset
# -> higher memory consumption
self.validator.validate(entity)
self.entity_validator.validate(entity)
write(entity)
self.validator.result()
self.logger.info(
f"{self.n_rows} records have been imported from {self.source_app}")
min_rows = self.dataset.get("min_rows", 1)
if self.mode == ImportMode.FULL and self.n_rows < min_rows:
# Default requirement for full imports is a non-empty dataset
self.logger.error(
f"Too few records imported: {self.n_rows} < {min_rows}")
def test_query(self):
reader = Reader(self.source, self.app, self.dataset())
reader._protect_row = lambda x: 'protected(' + x + ')'
reader.secure_attributes = []
query = iter(['a', 'b'])
self.assertEqual(['a', 'b'], list(reader._maybe_protect_rows(query)))
query = iter(['a', 'b'])
reader.secure_attributes = ['not', 'empty']
self.assertEqual([
'protected(a)',
'protected(b)',
], list(reader._maybe_protect_rows(query)))
def test_read(self):
reader = Reader({'query': ['a', 'b', 'c']}, self.app, self.dataset())
reader.datastore = mock.MagicMock()
reader._maybe_protect_rows = mock.MagicMock()
query_kwargs = {
'arraysize': 2000,
'name': 'import_cursor',
'withhold': True
}
result = reader.read()
self.assertEqual(reader._maybe_protect_rows.return_value, result)
reader._maybe_protect_rows.assert_called_with(
reader.datastore.query.return_value)
reader.datastore.query.assert_called_with('a\nb\nc', **query_kwargs)
reader.source = {'query': ['a', 'b', 'c'], 'recent': ['d', 'e']}
reader.mode = ImportMode.RECENT
reader.read()
reader.datastore.query.assert_called_with('a\nb\nc\nd\ne',
**query_kwargs)