def test_extraction_with_empty_query_result(self) -> None: """ Test Extraction with empty result from query """ with patch.object(SQLAlchemyExtractor, '_get_connection'): extractor = AthenaMetadataExtractor() extractor.init(self.conf) results = extractor.extract() self.assertEqual(results, None)
def test_extraction_with_single_result(self): # type: () -> None with patch.object(SQLAlchemyExtractor, '_get_connection') as mock_connection: connection = MagicMock() mock_connection.return_value = connection sql_execute = MagicMock() connection.execute = sql_execute table = {'schema_name': 'test_schema', 'name': 'test_table', 'description': '', 'cluster': self.conf['extractor.athena_metadata.{}'.format(AthenaMetadataExtractor.CATALOG_KEY)], } sql_execute.return_value = [ self._union( {'col_name': 'col_id1', 'col_type': 'bigint', 'col_description': 'description of id1', 'col_sort_order': 0, 'extras': None}, table), self._union( {'col_name': 'col_id2', 'col_type': 'bigint', 'col_description': 'description of id2', 'col_sort_order': 1, 'extras': None}, table), self._union( {'col_name': 'is_active', 'col_type': 'boolean', 'col_description': None, 'col_sort_order': 2, 'extras': None}, table), self._union( {'col_name': 'source', 'col_type': 'varchar', 'col_description': 'description of source', 'col_sort_order': 3, 'extras': None}, table), self._union( {'col_name': 'etl_created_at', 'col_type': 'timestamp', 'col_description': None, 'col_sort_order': 4, 'extras': 'partition key'}, table), self._union( {'col_name': 'ds', 'col_type': 'varchar', 'col_description': None, 'col_sort_order': 5, 'extras': None}, table) ] extractor = AthenaMetadataExtractor() extractor.init(self.conf) actual = extractor.extract() expected = TableMetadata('athena', self.conf['extractor.athena_metadata.{}'. format(AthenaMetadataExtractor.CATALOG_KEY)], 'test_schema', 'test_table', '', [ColumnMetadata('col_id1', 'description of id1', 'bigint', 0), ColumnMetadata('col_id2', 'description of id2', 'bigint', 1), ColumnMetadata('is_active', None, 'boolean', 2), ColumnMetadata('source', 'description of source', 'varchar', 3), ColumnMetadata('etl_created_at', 'partition key', 'timestamp', 4), ColumnMetadata('ds', None, 'varchar', 5)]) self.assertEqual(expected.__repr__(), actual.__repr__()) self.assertIsNone(extractor.extract())