def test_main_happy_path_s3(self, logger_setup, es_conn, bulk, local_time):
logger = Mock()
logger_setup.return_value = logger
es = Mock()
es.indices.exists_alias.return_value = False
es_conn.return_value = es
bulk.side_effect = self.capture_actions
# GMT: Monday, September 9, 2019 4:00:00 AM
# EDT: Monday, September 9, 2019 12:00:00 AM
local_time.return_value = 1568001600
self.optional[-1] = toAbsolute('__fixtures__/from_s3.ndjson')
self.optional.append('--metadata')
self.optional.append(toAbsolute('__fixtures__/metadata.json'))
argv = build_argv(self.optional)
with captured_output(argv) as (out, err):
sut.main()
# Expected index create calls
es.indices.create.assert_any_call(index='onion-v1', ignore=400)
es.indices.create.assert_any_call(index='onion-v2', ignore=400)
es.indices.create.assert_any_call(index='onion-v1', body=ANY)
# Expected index put_alias calls
es.indices.put_alias.assert_called_once_with(name='onion',
index='onion-v1')
# Expected index delete calls
es.indices.delete.assert_called_once_with(index='onion-v1')
# Bulk
bulk.assert_called_once_with(es,
actions=ANY,
index='onion-v1',
doc_type='complaint',
chunk_size=20000,
refresh=True)
self.validate_actions(toAbsolute('__fixtures__/exp_s3.ndjson'))
logger.info.assert_any_call('Deleting and recreating onion-v1')
logger.info.assert_any_call(
'Loading data into onion-v1 with doc_type complaint')
logger.info.assert_any_call('chunk retrieved, now bulk load')
logger.info.assert_any_call('1,001 records indexed, total = 1,001')
logger.info.assert_any_call('Adding alias onion for index onion-v1')
def test_main_happy_path_socrata(self, logger_setup, es_conn, bulk):
logger = Mock()
logger_setup.return_value = logger
es = Mock()
es.indices.exists_alias.return_value = False
es_conn.return_value = es
bulk.side_effect = self.capture_actions
self.optional.insert(0, '--dump-config')
self.optional[-1] = toAbsolute('../../tests/__fixtures__/ccdb.ndjson')
argv = build_argv(self.optional)
with captured_output(argv) as (out, err):
sut.main()
# Expected index create calls
es.indices.create.assert_any_call(index='onion-v1', ignore=400)
es.indices.create.assert_any_call(index='onion-v2', ignore=400)
es.indices.create.assert_any_call(index='onion-v1', body=ANY)
# Expected index put_alias calls
es.indices.put_alias.assert_called_once_with(name='onion',
index='onion-v1')
# Expected index delete calls
es.indices.delete.assert_called_once_with(index='onion-v1')
# Bulk
bulk.assert_called_once_with(es,
actions=ANY,
index='onion-v1',
doc_type='complaint',
chunk_size=20000,
refresh=True)
self.validate_actions(toAbsolute('__fixtures__/exp_socrata.ndjson'))
logger.info.assert_any_call('Running index_ccdb with')
logger.info.assert_any_call('Deleting and recreating onion-v1')
logger.info.assert_any_call(
'Loading data into onion-v1 with doc_type complaint')
logger.info.assert_any_call('chunk retrieved, now bulk load')
logger.info.assert_any_call('1,001 records indexed, total = 1,001')
logger.info.assert_any_call('Adding alias onion for index onion-v1')
def test_main_transport_error(self, logger_setup, es_conn, bulk):
from elasticsearch import TransportError
logger = Mock()
logger_setup.return_value = logger
es = Mock()
es.indices.exists_alias.return_value = False
es_conn.return_value = es
bulk.side_effect = TransportError(404, 'oops')
argv = build_argv(self.optional)
with captured_output(argv) as (out, err):
with self.assertRaises(SystemExit):
sut.main()
# Rollback
es.indices.put_alias.assert_called_once_with(name='onion',
index='onion-v2')
self.assertEqual(logger.error.call_count, 1)