def test_parquet_with_delimited_hint_hint_first(self, mock_Session): mock_source = Mock(name='source') mock_target = Mock(name='target') @contextmanager def mysource(records_format=None): self.assertIsInstance(records_format, ParquetRecordsFormat) yield mock_source @contextmanager def mytarget(db_engine): yield mock_target mock_job_name = Mock(name='job_name') mock_session = mock_Session.return_value mock_records = mock_session.records mock_records.sources.mysource = mysource mock_records.targets.mytarget = mytarget config = { 'fail_if_dont_understand': False, 'source': { 'compression': 'GZIP', 'format': 'parquet', }, 'target': { 'db_name': 'foo' }, } with self.assertRaises(NotImplementedError): run_records_mover_job(source_method_name='mysource', target_method_name='mytarget', job_name=mock_job_name, config=config)
def test_variant_with_no_hints(self, mock_Session): mock_source = Mock(name='source') mock_target = Mock(name='target') @contextmanager def mysource(initial_hints=None, records_format=None): self.assertIsNone(initial_hints) self.assertEqual(records_format.variant, 'vertica') yield mock_source @contextmanager def mytarget(db_engine): yield mock_target mock_job_name = Mock(name='job_name') mock_session = mock_Session.return_value mock_records = mock_session.records mock_records.sources.mysource = mysource mock_records.targets.mytarget = mytarget config = { 'fail_if_dont_understand': False, 'source': { 'variant': 'vertica', }, 'target': { 'db_name': 'foo' }, } out = run_records_mover_job(source_method_name='mysource', target_method_name='mytarget', job_name=mock_job_name, config=config) mock_records.move.assert_called() self.assertEqual(out, mock_records.move.return_value)
def test_parquet(self, mock_Session): mock_source = Mock(name='source') mock_target = Mock(name='target') @contextmanager def mysource(records_format=None): self.assertIsInstance(records_format, ParquetRecordsFormat) yield mock_source @contextmanager def mytarget(db_engine): yield mock_target mock_job_name = Mock(name='job_name') mock_session = mock_Session.return_value mock_records = mock_session.records mock_records.sources.mysource = mysource mock_records.targets.mytarget = mytarget config = { 'fail_if_dont_understand': False, 'source': { 'format': 'parquet' }, 'target': { 'db_name': 'foo' }, } out = run_records_mover_job(source_method_name='mysource', target_method_name='mytarget', job_name=mock_job_name, config=config) mock_records.move.assert_called() self.assertEqual(out, mock_records.move.return_value)
def test_run_records_mover_job(self, mock_Session): mock_source = Mock(name='source') mock_target = Mock(name='target') def mysource(a: int, schema_name: str, google_cloud_creds, existing_table_handling=None, spectrum_base_url=None): self.assertEqual(a, 1) self.assertEqual(google_cloud_creds, mock_session.creds.google_sheets.return_value) self.assertEqual(schema_name, 'myschema') self.assertEqual(spectrum_base_url, 'spectrumschema') return mock_source def mytarget(b: int, db_engine): assert b == 2 self.assertEqual(db_engine, mock_session.get_db_engine.return_value) return mock_target mock_job_name = Mock(name='job_name') mock_session = mock_Session.return_value mock_db_facts = { 'redshift_spectrum_base_url_myschema': 'spectrumschema' } mock_session.creds.db_facts.return_value = mock_db_facts mock_records = mock_session.records mock_records.sources.mysource = mysource mock_records.targets.mytarget = mytarget config = { 'fail_if_dont_understand': False, 'source': { 'a': 1, 'db_name': 'mydbname', 'gcp_creds_name': 'mygcpcreds', 'existing_table': 'drop_and_recreate', 'schema_name': 'myschema' }, 'target': { 'b': 2, 'db_name': 'foo', } } out = run_records_mover_job(source_method_name='mysource', target_method_name='mytarget', job_name=mock_job_name, config=config) mock_records.move.assert_called() self.assertEqual(out, mock_records.move.return_value)
def test_hints_with_no_variant(self, mock_Session): mock_source = Mock(name='source') mock_target = Mock(name='target') @contextmanager def mysource(initial_hints=None, records_format=None): self.assertTrue(initial_hints is not None) self.assertEqual(',', initial_hints['field-delimiter']) self.assertEqual('GZIP', initial_hints['compression']) self.assertEqual('all', initial_hints['quoting']) self.assertIsNone(records_format) yield mock_source @contextmanager def mytarget(db_engine): yield mock_target mock_job_name = Mock(name='job_name') mock_session = mock_Session.return_value mock_records = mock_session.records mock_records.sources.mysource = mysource mock_records.targets.mytarget = mytarget config = { 'fail_if_dont_understand': False, 'source': { 'field-delimiter': ',', 'compression': 'GZIP', 'quoting': 'all', }, 'target': { 'db_name': 'foo' } } out = run_records_mover_job(source_method_name='mysource', target_method_name='mytarget', job_name=mock_job_name, config=config) mock_records.move.assert_called() self.assertEqual(out, mock_records.move.return_value)
def test_run_records_mover_hint_and_then_variant(self, mock_Session): mock_source = Mock(name='source') mock_target = Mock(name='target') def mysource(): return mock_source def mytarget(records_format: Optional[BaseRecordsFormat] = None): assert isinstance(records_format, DelimitedRecordsFormat) self.assertEqual(records_format.variant, 'csv') self.assertIsNone(records_format.hints['escape']) return mock_target mock_job_name = Mock(name='job_name') mock_session = mock_Session.return_value mock_db_facts = { 'redshift_spectrum_base_url_myschema': 'spectrumschema' } mock_session.creds.db_facts.return_value = mock_db_facts mock_records = mock_session.records mock_records.sources.mysource = mysource mock_records.targets.mytarget = mytarget config = { 'fail_if_dont_understand': False, 'source': {}, 'target': { 'compression': None, 'variant': 'csv' } } out = run_records_mover_job(source_method_name='mysource', target_method_name='mytarget', job_name=mock_job_name, config=config) mock_records.move.assert_called() self.assertEqual(out, mock_records.move.return_value)